fix syntax issue

Author: nitins17

.github/workflows/wheel_tests_continuous.yml

@@ -22,14 +22,14 @@ on:
     - cron: "0 */2 * * *" # Run once every 2 hours
   pull_request:
     branches:
-      - main  
+      - main
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
   cancel-in-progress: ${{ !contains(github.ref, 'release/') && github.ref != 'main' }}
 
 jobs:
-  build_jaxlib_artifact:
+  build-jaxlib-artifact:
     uses: ./.github/workflows/build_artifacts.yml
     strategy:
         fail-fast: false # don't cancel all jobs on failure
@@ -46,12 +46,12 @@ jobs:
         upload_artifacts_to_gcs: true
         gcs_upload_uri: 'gs://general-ml-ci-transient/jax-github-actions/jax-fork/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
 
-  build_cuda_artifacts:
-    uses: ./.github/workflows/build_artifacts.yml      
+  build-cuda-artifacts:
+    uses: ./.github/workflows/build_artifacts.yml
     strategy:
         fail-fast: false # don't cancel all jobs on failure
         matrix:
-          # Python values need to match the matrix stategy in the GPU tests job below
+          # Python values need to match the matrix stategy in the CUDA tests job below
           runner: ["linux-x86-n2-16"]
           artifact: ["jax-cuda-plugin", "jax-cuda-pjrt"]
           python: ["3.10",]
@@ -62,48 +62,54 @@ jobs:
       clone_main_xla: 1
       upload_artifacts_to_gcs: true
       gcs_upload_uri: 'gs://general-ml-ci-transient/jax-github-actions/jax-fork/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
- 
-  run_pytest_cpu:
-    needs: build_jaxlib_artifact
-    uses: ./.github/workflows/pytest_cpu.yml      
+
+  run-pytest-cpu:
+    needs: build-jaxlib-artifact
+    uses: ./.github/workflows/pytest_cpu.yml
     strategy:
         fail-fast: false # don't cancel all jobs on failure
         matrix:
-          # Runner OS and Python values need to match the matrix stategy in the build_jaxlib_artifact job above
+          # Runner OS and Python values need to match the matrix stategy in the
+          # build_jaxlib_artifact job above
           runner: ["linux-x86-n2-64", "linux-arm64-t2a-48", "windows-x86-n2-64"]
           python: ["3.10",]
           enable-x64: [1, 0]
     with:
       runner: ${{ matrix.runner }}
       python: ${{ matrix.python }}
       enable-x64:  ${{ matrix.enable-x64 }}
-      gcs_download_uri: ${{ needs.build_jaxlib_artifact.outputs.gcs_upload_uri }}
-      install_jax_current_commit: "1"
-
+      gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri }}
 
-  run_pytest_gpu:
-    needs: [build_jaxlib_artifact, build_cuda_artifacts]
-    uses: ./.github/workflows/pytest_cuda.yml      
+  run-pytest-cuda:
+    needs: [build-jaxlib-artifact, build-cuda-artifacts]
+    uses: ./.github/workflows/pytest_cuda.yml
     strategy:
         fail-fast: false # don't cancel all jobs on failure
         matrix:
-          # Python values need to match the matrix stategy in the build artifacts job above
-          runner: ["linux-x86-g2-48-l4-4gpu",]
+          # Python values need to match the matrix stategy in the artifact build jobs above
+          runner: ["linux-x86-g2-48-l4-4gpu", "linux-x86-a3-8g-h100-8gpu"]
           python: ["3.10",]
           cuda: ["12.3", "12.1"]
           enable-x64: [1, 0]
+          exclude:
+            # Run only a single configuration on H100 to save resources
+            - runner: "linux-x86-a3-8g-h100-8gpu"
+              python: "3.10"
+              cuda: "12.1"
+            - runner: "linux-x86-a3-8g-h100-8gpu"
+              python: "3.10"
+              enable-x64: 0
     with:
       runner: ${{ matrix.runner }}
       python: ${{ matrix.python }}
       cuda:  ${{ matrix.cuda }}
       enable-x64:  ${{ matrix.enable-x64 }}
-       # GCS upload URI is the same for both artifact build jobs
-      gcs_download_uri: ${{ needs.build_jaxlib_artifact.outputs.gcs_upload_uri }}
-      install_jax_current_commit: "1"
+      # GCS upload URI is the same for both artifact build jobs
+      gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri }}
 
-  run_bazel_test_gpu:
-    needs: [build_jaxlib_artifact, build_cuda_artifacts]
-    uses: ./.github/workflows/bazel_cuda_non_rbe.yml      
+  run-bazel-test-cuda:
+    needs: [build-jaxlib-artifact, build-cuda-artifacts]
+    uses: ./.github/workflows/bazel_cuda_non_rbe.yml
     strategy:
         fail-fast: false # don't cancel all jobs on failure
         matrix:
@@ -115,4 +121,5 @@ jobs:
       runner: ${{ matrix.runner }}
       python: ${{ matrix.python }}
       enable-x64:  ${{ matrix.enable-x64 }}
-      gcs_download_uri: ${{ needs.build_jaxlib_artifact.outputs.gcs_upload_uri }}
+      # GCS upload URI is the same for both artifact build jobs
+      gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri }}