Install test dependencies from test-requirements.txt instead of requirements.in

PiperOrigin-RevId: 736878834

Author: nitins17

.github/workflows/pytest_cpu.yml

@@ -116,21 +116,8 @@ jobs:
           exit 1
       - name: Install Python dependencies
         run: |
-          # Remove installation of NVIDIA wheels for CPU tests.
-          sed -i 's/-r gpu-test-requirements.txt/# -r gpu-test-requirements.txt/g' build/requirements.in
-
-          # TODO(srnitin): Remove after uv is installed in the Windows Dockerfile
           $JAXCI_PYTHON -m pip install uv~=0.5.30
-          # python 3.13t cannot compile zstandard 0.23.0 due to
-          # https://github.com/indygreg/python-zstandard/issues/231. Remove this once zstandard
-          # has a prebuilt wheel for 3.13t or an env marker is available for free threading python
-          # in requirements.in.
-          if [[ $JAXCI_PYTHON =~ "python3.13-nogil" ]]; then
-            grep -v "zstandard" build/requirements.in > build/requirements_without_zstandard.txt
-            $JAXCI_PYTHON -m uv pip install -r build/requirements_without_zstandard.txt
-          else
-            $JAXCI_PYTHON -m uv pip install -r build/requirements.in
-          fi
+          $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

.github/workflows/pytest_cuda.yml

@@ -100,17 +100,7 @@ jobs:
           echo "Skipping the test run."
           exit 1
       - name: Install Python dependencies
-        run: |
-          # python 3.13t cannot compile zstandard 0.23.0 due to
-          # https://github.com/indygreg/python-zstandard/issues/231. Remove this once zstandard
-          # has a prebuilt wheel for 3.13t or an env marker is available for free threading python
-          # in requirements.in.
-          if [[ $JAXCI_PYTHON =~ "python3.13-nogil" ]]; then
-            grep -v "zstandard" build/requirements.in > build/requirements_without_zstandard.txt
-            $JAXCI_PYTHON -m uv pip install -r build/requirements_without_zstandard.txt
-          else
-            $JAXCI_PYTHON -m uv pip install -r build/requirements.in
-          fi
+        run: $JAXCI_PYTHON -m uv pip install -r build/test-requirements.txt
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

build/test-requirements.txt

@@ -7,7 +7,7 @@ flatbuffers
 hypothesis
 mpmath>=1.3
 pillow>=10.4.0
-# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t 
+# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t
 portpicker; python_version<"3.13"
 pytest-xdist
 wheel
@@ -19,3 +19,6 @@ 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"