Fixed failed test due to uint8 overflow

tests/compile/function/test_function.py

@@ -11,6 +11,7 @@ from pytensor.compile import shared
 from pytensor.compile.function import function, function_dump
 from pytensor.compile.io import In
 from pytensor.configdefaults import config
+from pytensor.npy_2_compat import UintOverflowError
 from pytensor.tensor.type import (
     bscalar,
     bvector,
@@ -166,12 +167,12 @@ class TestFunctionIn:
         # Value too big for a, silently ignored
         assert np.array_equal(f([2**20], np.ones(1, dtype="int8"), 1), [2])
 
-        # Value too big for b, raises TypeError
-        with pytest.raises(TypeError):
+        # Value too big for b, raises OverflowError (in numpy >= 2.0... TypeError in numpy < 2.0)
+        with pytest.raises(UintOverflowError):
             f([3], [312], 1)
 
-        # Value too big for c, raises TypeError
-        with pytest.raises(TypeError):
+        # Value too big for c, raises OverflowError
+        with pytest.raises(UintOverflowError):
             f([3], [6], 806)
 
     def test_in_allow_downcast_floatX(self):

tests/compile/function/test_pfunc.py

@@ -9,6 +9,7 @@ from pytensor.compile.io import In
 from pytensor.compile.sharedvalue import shared
 from pytensor.configdefaults import config
 from pytensor.graph.utils import MissingInputError
+from pytensor.npy_2_compat import UintOverflowError
 from pytensor.tensor.math import sum as pt_sum
 from pytensor.tensor.type import (
     bscalar,
@@ -237,12 +238,12 @@ class TestPfunc:
         # Value too big for a, silently ignored
         assert np.all(f([2**20], np.ones(1, dtype="int8"), 1) == 2)
 
-        # Value too big for b, raises TypeError
-        with pytest.raises(TypeError):
+        # Value too big for b, raises OverflowError in numpy >= 2.0, TypeError in numpy <2.0
+        with pytest.raises(UintOverflowError):
             f([3], [312], 1)
 
-        # Value too big for c, raises TypeError
-        with pytest.raises(TypeError):
+        # Value too big for c, raises OverflowError in numpy >= 2.0, TypeError in numpy <2.0
+        with pytest.raises(UintOverflowError):
             f([3], [6], 806)
 
     def test_param_allow_downcast_floatX(self):
@@ -327,8 +328,8 @@ class TestPfunc:
         with pytest.raises(TypeError):
             g([3], np.array([6], dtype="int16"), 0)
 
-        # Value too big for b, raises TypeError
-        with pytest.raises(TypeError):
+        # Value too big for b, raises OverflowError in numpy >= 2.0, TypeError in numpy <2.0
+        with pytest.raises(UintOverflowError):
             g([3], [312], 0)
 
         h = pfunc([a, b, c], (a + b + c))  # Default: allow_input_downcast=None
@@ -336,7 +337,9 @@ class TestPfunc:
         assert np.all(h([3], [6], 0) == 9)
         with pytest.raises(TypeError):
             h([3], np.array([6], dtype="int16"), 0)
-        with pytest.raises(TypeError):
+
+        # Value too big for b, raises OverflowError in numpy >= 2.0, TypeError in numpy <2.0
+        with pytest.raises(UintOverflowError):
             h([3], [312], 0)
 
     def test_allow_downcast_floatX(self):

tests/tensor/test_elemwise.py

@@ -40,7 +40,27 @@ from pytensor.tensor.type import (
 )
 from tests import unittest_tools
 from tests.link.test_link import make_function
-from tests.tensor.test_math import reduce_bitwise_and
+
+
+def reduce_bitwise_and(x, axis=-1, dtype="int8"):
+    """Helper function for TestCAReduce"""
+    if dtype == "uint8":
+        # in numpy version >= 2.0, out of bounds uint8 values are not converted
+        identity = np.array((255,), dtype=dtype)[0]
+    else:
+        identity = np.array((-1,), dtype=dtype)[0]
+
+    shape_without_axis = tuple(s for i, s in enumerate(x.shape) if i != axis)
+    if 0 in shape_without_axis:
+        return np.empty(shape=shape_without_axis, dtype=x.dtype)
+
+    def custom_reduce(a):
+        out = identity
+        for i in range(a.size):
+            out = np.bitwise_and(a[i], out)
+        return out
+
+    return np.apply_along_axis(custom_reduce, axis, x)
 
 
 class TestDimShuffle(unittest_tools.InferShapeTester):

tests/tensor/test_math.py

@@ -3444,22 +3444,6 @@ class TestSumMeanMaxMinArgMaxVarReduceAxes:
             x.var(a)
 
 
-def reduce_bitwise_and(x, axis=-1, dtype="int8"):
-    identity = np.array((-1,), dtype=dtype)[0]
-
-    shape_without_axis = tuple(s for i, s in enumerate(x.shape) if i != axis)
-    if 0 in shape_without_axis:
-        return np.empty(shape=shape_without_axis, dtype=x.dtype)
-
-    def custom_reduce(a):
-        out = identity
-        for i in range(a.size):
-            out = np.bitwise_and(a[i], out)
-        return out
-
-    return np.apply_along_axis(custom_reduce, axis, x)
-
-
 def test_clip_grad():
     # test the gradient of clip
     def func(x, y, z):