jax

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