Experiment with resuable workflows

Author: nitins17

.github/workflows/build_artifacts.yml

@@ -4,7 +4,7 @@ on:
   pull_request:
     branches:
       - main
-  workflow_dispatch:
+  workflow_call:
 
 jobs:
   build:

.github/workflows/pytest_cpu_reuse.yml

@@ -0,0 +1,62 @@
+name: Run Pytest CPU tests (resuable workflow)
+
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  build_jaxlib_artifacts:
+    strategy:
+      matrix:
+        runner: ["windows-x86-n2-64", "linux-x86-n2-64", "linux-arm64-t2a-48"]
+        python: ["3.10"]
+        exclude:
+          # Only build the jaxlib artifact.
+          - artifact: "jax"
+          - artifact: "jax-cuda-plugin"
+          - artifact: "jax-cuda-pjrt"
+
+    uses: ./.github/workflows/build_artifacts.yml
+
+  run_pytest:
+    continue-on-error: true
+    defaults:
+      run:
+        # Explicitly set the shell to bash to override the default Windows environment, i.e, cmd.
+        shell: bash
+    strategy:
+      matrix:
+        runner: ["windows-x86-n2-64", "linux-x86-n2-64", "linux-arm64-t2a-48"]
+        python: ["3.10"]
+
+    runs-on: ${{ matrix.runner }}
+    container: ${{ (contains(matrix.runner, 'linux-x86') && 'gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest') ||
+                (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest') ||
+                (contains(matrix.runner, 'windows-x86') && null) }}
+
+    env:
+      # Disable running `run_docker_container.sh`. JAX self-hosted runners already run in a Docker
+      # container.
+      JAXCI_RUN_DOCKER_CONTAINER: 0
+      # Use RBE to build the artifacts where possible (Linux x86 and Windows).
+      JAXCI_BUILD_ARTIFACT_WITH_RBE: 1
+      # Setup the test environment (disable x64 mode and clone XLA at HEAD)
+      JAXCI_SETUP_TEST_ENVIRONMENT: 1
+      JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
+
+    steps:
+      - uses: actions/checkout@v3
+      # Halt for testing
+      - name: Wait For Connection
+        uses: ./actions/ci_connection/
+      - name: Install pytest
+        env:
+          JAXCI_PYTHON: python${{ matrix.python }}
+        run: $JAXCI_PYTHON -m pip install pytest
+      - name: Install dependencies
+        env:
+          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"