diff --git a/.github/workflows/rocm-ci.yml b/.github/workflows/rocm-ci.yml
index aec436677..3eda5116d 100644
--- a/.github/workflows/rocm-ci.yml
+++ b/.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"
diff --git a/build/rocm/ci_build b/build/rocm/ci_build
index 469e9434d..b492c808d 100755
--- a/build/rocm/ci_build
+++ b/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)
 
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index 8ee9b22ff..178be9bff 100644
--- a/third_party/xla/workspace.bzl
+++ b/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