Fix jax wheel pattern in pytest workflows

The jax* pattern incorrectly installs jax_cuda_pjrt wheels when installing jax wheels.

Author: kanglant

.github/workflows/pytest_cpu.yml

@@ -101,7 +101,7 @@ jobs:
 
           # Download the "jax" wheel from GCS if inputs.install-jax-current-commit is not set to 1
           if [[ "${{ inputs.install-jax-current-commit }}" != 1 ]]; then
-            gsutil -m cp -r "${{ inputs.gcs_download_uri }}"/jax*py3*none*.whl $(pwd)/dist/
+            gsutil -m cp -r "${{ inputs.gcs_download_uri }}"/jax-*py3*none*.whl $(pwd)/dist/
           fi
       - name: Download jaxlib wheel from GCS (Windows runs)
         id: download-wheel-artifacts-w
@@ -117,7 +117,7 @@ jobs:
 
           # Download the "jax" wheel from GCS if inputs.install-jax-current-commit is not set to 1
           if not "${{ inputs.install-jax-current-commit }}"=="1" (
-            gsutil -m cp -r "${{ inputs.gcs_download_uri }}"/jax*py3*none*.whl dist/
+            gsutil -m cp -r "${{ inputs.gcs_download_uri }}"/jax-*py3*none*.whl dist/
           )
       - name: Skip the test run if the wheel artifacts were not downloaded successfully
         if: steps.download-wheel-artifacts-nw.outcome == 'failure' || steps.download-wheel-artifacts-w.outcome == 'failure'

.github/workflows/pytest_cuda.yml

@@ -100,7 +100,7 @@ jobs:
 
           # Download the "jax" wheel from GCS if inputs.install-jax-current-commit is not set to 1
           if [[ "${{ inputs.install-jax-current-commit }}" != 1 ]]; then
-            gsutil -m cp -r "${{ inputs.gcs_download_uri }}"/jax*py3*none*.whl $(pwd)/dist/
+            gsutil -m cp -r "${{ inputs.gcs_download_uri }}"/jax-*py3*none*.whl $(pwd)/dist/
           fi
       - name: Skip the test run if the wheel artifacts were not downloaded successfully
         if: steps.download-wheel-artifacts.outcome == 'failure'