Update benchmarks.yml

Author: juliagmt-google

.github/workflows/benchmarks.yml

@@ -29,27 +29,38 @@ jobs:
       JAXCI_HERMETIC_PYTHON_VERSION: 3.11
 
     steps:
-      - uses: actions/checkout@v3
-      # - name: Build jaxlib
-      #   env:
-      #     JAXCI_CLONE_MAIN_XLA: 1
-      #   run: ./ci/build_artifacts.sh "jaxlib"
-      # - name: Build jax-cuda-plugin
-      #   env:
-      #     JAXCI_CLONE_MAIN_XLA: 1
-      #   run: ./ci/build_artifacts.sh "jax-cuda-plugin"
-      # - name: Build jax-cuda-pjrt
-      #   env:
-      #     JAXCI_CLONE_MAIN_XLA: 1
-      #   run: ./ci/build_artifacts.sh "jax-cuda-pjrt"
-      - name: Install dependencies
+      - name: Checkout JAX Fork
+        uses: actions/checkout@v3
+        with:
+          repository: 'google-ml-infra/jax-fork'
+          path: jax-fork
+
+      - name: Checkout XLA
+        uses: actions/checkout@v3
+        with:
+          repository: 'openxla/xla' # Or your XLA fork
+          path: xla
+
+      - name: Install JAX Dependencies
+        working-directory: jax-fork
         run: |
           python -m pip install --upgrade pip
           pip install pytest
           pip install absl-py
-          # pip install -U jax
-          pip install -U "jax[cuda12]"
+          pip install "jax[cuda12_pip]"  # Adjust CUDA version if needed
           pip install google-benchmark
-      - name: Run Multiprocess GPU Test
-        run: |
-          python -m pytest tests/multiprocess_gpu_test.py
+
+      - name: Run JAX Multiprocess GPU Test
+        working-directory: jax-fork
+        continue-on-error: true
+        run: python -m pytest tests/multiprocess_gpu_test.py
+      
+      - name: Build XLA GPU Atomic Test
+        working-directory: xla
+        continue-on-error: true
+        run: bazel build -c opt --config=cuda //xla/service/gpu/tests:gpu_atomic_test
+      
+      - name: Run XLA GPU Atomic Test
+        working-directory: xla
+        continue-on-error: true
+        run: bazel test -c opt --config=cuda //xla/service/gpu/tests:gpu_atomic_test