jax

bec3b790 - Add workflow for build wheels and running bazel tests on ROCm