Update _NPY_MAXDIMS (#7167)

NoureldinYosri · pavoljuhas · web-flow · commit 815649ca2f80 · 2025-03-25T07:27:48.000Z
* Update `_NPY_MAXDIMS` numpy 1.x supported a maximum of 32 dimensions and then raised an error .. this was an issue for large scale simulations with qubits up to 35 for state vector simulation and up to 18 qubits for density matrix simulations. As a workaround to support these cases we worked with 1D and 2D arrays for these cases which was very inefficient. now numpy 2.0 supports up to 64 dimensions * use the value from numpy instead of hardcoding it * use suggestion from @maffoo * Update cirq-core/cirq/linalg/transformations.py Co-authored-by: Pavol Juhas <juhas@google.com> * nit * Update transformations.py * Handle negative dimensions in can_numpy_support_shape And add a direct unit test. --------- Co-authored-by: Pavol Juhas <juhas@google.com>
diff --git a/cirq-core/cirq/linalg/transformations.py b/cirq-core/cirq/linalg/transformations.py
@@ -15,6 +15,7 @@
 """Utility methods for transforming matrices or vectors."""
 
 import dataclasses
+import functools
 from typing import Any, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
@@ -29,8 +30,6 @@
 # user provides a different np.array([]) value.
 RaiseValueErrorIfNotProvided: np.ndarray = np.array([])
 
-_NPY_MAXDIMS = 32  # Should be changed once numpy/numpy#5744 is resolved.
-
 
 def reflection_matrix_pow(reflection_matrix: np.ndarray, exponent: float):
     """Raises a matrix with two opposing eigenvalues to a power.
@@ -807,6 +806,15 @@ def transpose_flattened_array(t: np.ndarray, shape: Sequence[int], axes: Sequenc
     return ret
 
 
+@functools.cache
+def _can_numpy_support_dims(num_dims: int) -> bool:
+    try:
+        _ = np.empty((1,) * num_dims)
+        return True
+    except ValueError:  # pragma: no cover
+        return False
+
+
 def can_numpy_support_shape(shape: Sequence[int]) -> bool:
     """Returns whether numpy supports the given shape or not numpy/numpy#5744."""
-    return len(shape) <= _NPY_MAXDIMS
+    return min(shape, default=0) >= 0 and _can_numpy_support_dims(len(shape))
diff --git a/cirq-core/cirq/linalg/transformations_test.py b/cirq-core/cirq/linalg/transformations_test.py
@@ -648,3 +648,8 @@ def test_transpose_flattened_array(num_dimensions):
         assert np.array_equal(want, got)
         got = linalg.transpose_flattened_array(A.reshape(shape), shape, axes).reshape(want.shape)
         assert np.array_equal(want, got)
+
+
+@pytest.mark.parametrize('shape, result', [((), True), (30 * (1,), True), ((-3, 1, -1), False)])
+def test_can_numpy_support_shape(shape: tuple[int, ...], result: bool) -> None:
+    assert linalg.can_numpy_support_shape(shape) is result
