jax

cd77ba32 - use rocm_rbe also for build