Use bazel for PR tests (#216)

* Use bazel for running pre-merge CI tests

* Don't use HEREDOC

* Fix block text

* Use bash array

* Add bazel install

* Put Bazel in the build image

* Use Bazelisk

* Remove bazel install in Docker

* Go back to upstream XLA

* Remove bazel test command from workflow

* Move test command to build container

* Fix string format typos

.github/workflows/rocm-ci.yml

@@ -59,6 +59,9 @@ jobs:
           name: rocm_jax_r${{ env.ROCM_VERSION }}_py${{ env.PYTHON_VERSION  }}_id${{ github.run_id }}
           path: ./dist/*.whl
       - name: Run tests
+        env:
+          GPU_COUNT: "8"
+          GFX: "gfx90a"
         run: |
           cd $WORKSPACE_DIR
           python3 build/rocm/ci_build test $TEST_IMAGE --test-cmd "pytest tests/core_test.py"

build/rocm/ci_build

@@ -127,6 +127,31 @@ def dist_wheels(
         ]
     )
 
+    # Add command for unit tests
+    cmd.extend(
+        [
+            "&&",
+            "bazel",
+            "test",
+            "-k",
+            "--jobs=4",
+            "--test_verbose_timeout_warnings=true",
+            "--test_output=all",
+            "--test_summary=detailed",
+            "--local_test_jobs=1",
+            "--test_env=JAX_ACCELERATOR_COUNT=%i" % 4, 
+            "--test_env=JAX_SKIP_SLOW_TESTS=0",
+            "--verbose_failures=true",
+            "--config=rocm",
+            "--action_env=ROCM_PATH=/opt/rocm",
+            "--action_env=TF_ROCM_AMDGPU_TARGETS=%s" % "gfx90a",
+            "--test_tag_filters=-multiaccelerator",
+            "--test_env=XLA_PYTHON_CLIENT_ALLOCATOR=platform",
+            "--test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow",
+            "//tests:gpu_tests",
+        ]
+    )
+
     LOG.info("Running: %s", cmd)
     _ = subprocess.run(cmd, check=True)
 

third_party/xla/workspace.bzl

@@ -29,7 +29,7 @@ def repo():
         name = "xla",
         sha256 = XLA_SHA256,
         strip_prefix = "xla-{commit}".format(commit = XLA_COMMIT),
-        urls = tf_mirror_urls("https://github.com/rocm/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)),
+        urls = tf_mirror_urls("https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)),
     )
 
     # For development, one often wants to make changes to the TF repository as well