File einops-fix-tests-numpy-2-0-0.patch of Package python-einops
diff --git a/einops/array_api.py b/einops/array_api.py
index 0377a31..2d2ab3d 100644
--- a/einops/array_api.py
+++ b/einops/array_api.py
@@ -42,7 +42,10 @@ def rearrange(tensor: Tensor, pattern: str, **axes_lengths) -> Tensor:
def asnumpy(tensor: Tensor):
import numpy as np
- return np.from_dlpack(tensor)
+ if hasattr(tensor, '__dlpack__'):
+ return np.from_dlpack(tensor)
+ else:
+ return tensor
Shape = Tuple
diff --git a/tests/test_ops.py b/tests/test_ops.py
index 67cd4b4..d5f04f9 100644
--- a/tests/test_ops.py
+++ b/tests/test_ops.py
@@ -112,7 +112,10 @@ def test_ellipsis_ops_imperative():
def test_rearrange_array_api():
- import numpy.array_api as xp
+ if numpy.__version__ < "2.0.0":
+ import numpy.array_api as xp
+ else:
+ import numpy as xp
from einops import array_api as AA
x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
@@ -123,7 +126,10 @@ def test_rearrange_array_api():
def test_reduce_array_api():
- import numpy.array_api as xp
+ if numpy.__version__ < "2.0.0":
+ import numpy.array_api as xp
+ else:
+ import numpy as xp
from einops import array_api as AA
x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
@@ -537,7 +543,10 @@ def test_repeat_symbolic():
def test_repeat_array_api():
- import numpy.array_api as xp
+ if numpy.__version__ < "2.0.0":
+ import numpy.array_api as xp
+ else:
+ import numpy as xp
from einops import array_api as AA
x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5])
diff --git a/tests/test_packing.py b/tests/test_packing.py
index 07afd39..c2990ce 100644
--- a/tests/test_packing.py
+++ b/tests/test_packing.py
@@ -270,8 +270,11 @@ def test_pack_unpack_against_numpy():
def test_pack_unpack_array_api():
+ if np.__version__ < "2.0.0":
+ import numpy.array_api as xp
+ else:
+ import numpy as xp
from einops import array_api as AA
- import numpy.array_api as xp
for case in cases:
shape = case.shape