Update benchmarks.yml

Author: juliagmt-google

.github/workflows/benchmarks.yml

@@ -16,37 +16,37 @@ on:
         - 'no'
 
 jobs:
-  Tests:
-      strategy:
-        # Don't fail fast - want to see results for all builds even if one fails.
-        fail-fast: false
-        matrix:
-          job_info: [
-            {
-              os: "linux-x86-n2-16",
-              container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest",
-              pretty_name: "XLA Linux x86 CPU",
-            },
-            {
-              os: "linux-arm64-c4a-16",
-              container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest",
-              pretty_name: "XLA Linux ARM64 CPU",
-            }
-          ]
-      name: ${{ matrix.job_info.pretty_name }}
-      runs-on: ${{ matrix.job_info.os }}
-      container: ${{ matrix.job_info.container }}
-      defaults:
-        run:
-          shell: bash
-      timeout-minutes: 30
-      steps:
-        - name: Print machine specs
-          run: |
-            lscpu
-            free -h  # Memory information
-            df -h    # Disk space information
-            uname -a # Kernel information
+  # Tests:
+  #     strategy:
+  #       # Don't fail fast - want to see results for all builds even if one fails.
+  #       fail-fast: false
+  #       matrix:
+  #         job_info: [
+  #           {
+  #             os: "linux-x86-n2-16",
+  #             container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest",
+  #             pretty_name: "XLA Linux x86 CPU",
+  #           },
+  #           {
+  #             os: "linux-arm64-c4a-16",
+  #             container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest",
+  #             pretty_name: "XLA Linux ARM64 CPU",
+  #           }
+  #         ]
+  #     name: ${{ matrix.job_info.pretty_name }}
+  #     runs-on: ${{ matrix.job_info.os }}
+  #     container: ${{ matrix.job_info.container }}
+  #     defaults:
+  #       run:
+  #         shell: bash
+  #     timeout-minutes: 30
+  #     steps:
+  #       - name: Print machine specs
+  #         run: |
+  #           lscpu
+  #           free -h  # Memory information
+  #           df -h    # Disk space information
+  #           uname -a # Kernel information
 
         # - name: Check Python Version in Container
         #   run: python3 --version
@@ -61,11 +61,11 @@ jobs:
         #       echo "Python 3 found."
         #     fi
 
-        - name: Checkout OpenXLA
-          uses: actions/checkout@v3
-          with:
-            repository: openxla/xla
-            ref: dev  # Specify the dev branch
+        # - name: Checkout OpenXLA
+        #   uses: actions/checkout@v3
+        #   with:
+        #     repository: openxla/xla
+        #     ref: dev  # Specify the dev branch
 
         # - name: Create results directory
         #   working-directory: xla
@@ -133,22 +133,22 @@ jobs:
         #     name: cpu-xla-benchmarks-${{ matrix.job_info.os }}
         #     path: xla/results
 
-  # build-xla-gpu-and-test:
-  #   runs-on: "linux-x86-n2-16" #linux-x86-g2-48-l4-4gpu # Use a GPU-enabled runner
-  #   container:
-  #     image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
-  #     options: --gpus all --privileged  # Might need privileged mode, use with caution
+  build-xla-gpu-and-test:
+    runs-on: "linux-x86-n2-16" #linux-x86-g2-48-l4-4gpu # Use a GPU-enabled runner
+    container:
+      image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
+      options: --gpus all --privileged  # Might need privileged mode, use with caution
 
-  #   steps:
-  #     - name: Checkout XLA
-  #       uses: actions/checkout@v3
-  #       with:
-  #         repository: openxla/xla  # Replace with your fork if needed
-  #     # - name: Checkout repository
-  #     #   uses: actions/checkout@v3
-  #     #   with:
-  #     #     repository: juliagmt-google/xla  
-  #     #     path: xla
+    # steps:
+    #   - name: Checkout XLA
+    #     uses: actions/checkout@v3
+    #     with:
+    #       repository: openxla/xla  # Replace with your fork if needed
+      - name: Checkout repository
+        uses: actions/checkout@v3
+        with:
+          repository: juliagmt-google/xla  
+          ref: dev
 
       - name: Print machine specs
         run: |