jax

Enable additional single-GPU bazel tests in ROCm CI

#816

Merged