From 031614c22b8aa8c60652d379e54baea187c427a0 Mon Sep 17 00:00:00 2001
From: Nitin Srinivasan <srnitin@google.com>
Date: Mon, 17 Mar 2025 08:58:18 -0700
Subject: [PATCH] Pin numpy~=2.1.0 in workflow file instead of
 test-requirements.txt

PiperOrigin-RevId: 737632771
---
 .github/workflows/pytest_cpu.yml | 5 +++++
 build/test-requirements.txt      | 5 +----
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/pytest_cpu.yml b/.github/workflows/pytest_cpu.yml
index d59f5606a..137f49c6d 100644
--- a/.github/workflows/pytest_cpu.yml
+++ b/.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
diff --git a/build/test-requirements.txt b/build/test-requirements.txt
index 84cd01d82..f0b315771 100644
--- a/build/test-requirements.txt
+++ b/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
\ No newline at end of file