jax

dba3b37d - Let the unit tests use build.py for setting up Bazel commands for unit tests (#582)