Pin numpy~=2.1.0 in workflow file instead of test-requirements.txt

PiperOrigin-RevId: 737632771

Author: nitins17

.github/workflows/pytest_cpu.yml

@@ -118,6 +118,11 @@ jobs:
         run: |
           $JAXCI_PYTHON -m pip install uv~=0.5.30
           $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
+
+          # CPU Pytests crash with NumPy 2.2+ on Linux Aarch64; b/399168632
+          if [[ $OS == "linux" && $ARCH == "aarch64" ]]; then
+            $JAXCI_PYTHON -m uv pip install numpy~=2.1.0
+          fi
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

build/test-requirements.txt

@@ -18,7 +18,4 @@ setuptools
 matplotlib~=3.8.4; python_version=="3.10"
 matplotlib; python_version>="3.11"
 opt-einsum
-auditwheel
-
-# CPU Pytests crash with NumPy 2.2+ on Linux Aarch64; b/399168632
-numpy~=2.1.0; platform_system == "Linux" and platform_machine == "aarch64"
+auditwheel