split test commands into separate scripts

Author: nitins17

.github/workflows/bazel_cpu_rbe.yml

@@ -33,4 +33,4 @@ jobs:
       - name: Wait For Connection
         uses: ./actions/ci_connection/
       - name: Run Bazel CPU Tests
-        run: ./ci/run_bazel_test.sh "ci/envs/run_tests/bazel_cpu_rbe.env"
+        run: ./ci/run_bazel_test_cpu_rbe.sh

.github/workflows/bazel_gpu_non_rbe.yml

@@ -33,6 +33,4 @@ jobs:
       - name: Build jax-cuda-pjrt
         run: ./ci/build_artifacts.sh "ci/envs/build_artifacts/jax-cuda-pjrt.env"
       - name: Run Bazel GPU tests locally
-        run: ./ci/run_bazel_test.sh "ci/envs/run_tests/bazel_gpu_non_rbe.env"
-      - name: Test if step runs
-        run: echo "This step ran!"
+        run: ./ci/run_bazel_test_gpu_non_rbe.sh 

.github/workflows/bazel_gpu_rbe.yml

@@ -27,4 +27,4 @@ jobs:
       - name: Wait For Connection
         uses: ./actions/ci_connection/
       - name: Run Bazel GPU tests using RBE
-        run: ./ci/run_bazel_test.sh "ci/envs/run_tests/bazel_gpu_rbe.env"
+        run: ./ci/run_bazel_test_gpu_rbe.sh

.github/workflows/pytest_cpu.yml

@@ -44,4 +44,4 @@ jobs:
           JAXCI_PYTHON: python${{ matrix.python }}
         run: $JAXCI_PYTHON -m pip install --upgrade     numpy=="2.0.0" scipy=="1.13.1"     wheel build pytest-xdist absl-py opt-einsum colorama portpicker     matplotlib 'importlib_metadata>=4.6' hypothesis flatbuffers filelock
       - name: Run Pytest CPU tests
-        run: ./ci/run_pytest.sh "ci/envs/run_tests/pytest_cpu.env"
+        run: ./ci/run_pytest_cpu.sh

.github/workflows/pytest_gpu.yml

@@ -41,4 +41,4 @@ jobs:
           JAXCI_PYTHON: python${{ matrix.python }}
         run: $JAXCI_PYTHON -m pip install -r build/test-requirements.txt
       - name: Run Pytest GPU tests
-        run: ./ci/run_pytest.sh "ci/envs/run_tests/pytest_gpu.env"
+        run: ./ci/run_pytest_gpu.sh

.github/workflows/pytest_tpu.yml

@@ -43,5 +43,4 @@ jobs:
           JAXCI_PYTHON: python${{ matrix.python }}
         run: $JAXCI_PYTHON -m pip install --pre libtpu-nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
       - name: Run Pytest TPU tests
-        env:
-        run: ./ci/run_pytest.sh "ci/envs/run_tests/pytest_tpu.env"
+        run: ./ci/run_pytest_tpu.sh