Add jax test script (#406)

* Add jax test script

* Adjust tag filters

* Enable rbe for jax-ut

* Fix docker url

* Add rbe settings manually

* Fix rbe pool name

* Build the tests instead of running them as they are still red

* Restore platform

* Use xla path passed as argument

Author: alekstheod

build_tools/rocm/platform/linux_x64/BUILD

@@ -13,7 +13,6 @@ platform(
     exec_properties = {
         "container-image": "rocm/tensorflow-build@sha256:7cd444ac48657fee2f5087fbda7766266704d3f8fb2299f681952ae4eabed060",
         "OSFamily": "Linux",
-        "Pool": "linux_x64_large",
     },
 )
 
@@ -27,6 +26,6 @@ platform(
     exec_properties = {
         "container-image": "rocm/tensorflow-build@sha256:7cd444ac48657fee2f5087fbda7766266704d3f8fb2299f681952ae4eabed060",
         "OSFamily": "Linux",
-        "Pool": "amd_gpu",
+        "Pool": "linux_x64_gpu",
     },
 )

build_tools/rocm/run_jax_ut.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+set -e
+
+JAX_DIR=$1
+XLA_DIR=$2
+
+pushd $JAX_DIR
+
+python build/build.py build \
+    --wheels=jax-rocm-plugin \
+    --configure_only \
+    --local_xla_path=${XLA_DIR} \
+    --python_version=3.12
+
+# TODO: run the tests when they are green
+bazel build \
+    --config=rocm \
+    --build_tag_filters=cpu,gpu,-tpu,-config-cuda-only \
+    --test_tag_filters=cpu,gpu,-tpu,-config-cuda-only \
+    --action_env=TF_ROCM_AMDGPU_TARGETS=gfx908,gfx90a,gfx942 \
+    --//jax:build_jaxlib=true \
+    "//tests/..."
+
+popd

build_tools/rocm/run_xla_ci_build.sh

@@ -71,7 +71,6 @@ EXCLUDED_TESTS=(
   # //xla/tools/multihost_hlo_runner:functional_hlo_runner_test
   FunctionalHloRunnerTest.Sharded2DevicesHloUnoptimizedSnapshot
   FunctionalHloRunnerTest.ShardedComputationUnderStreamCapture
-
 )
 
 SCRIPT_DIR=$(realpath $(dirname $0))