Only require input_ndim and not input_broadcastable in DimShuffle

Author: ricardoV94

tests/tensor/rewriting/test_math.py

@@ -3749,7 +3749,7 @@ def test_local_log_sum_exp_maximum():
     check_max_log_sum_exp(x, axis=(0, 1, 2), dimshuffle_op=None)
 
     # If a transpose is applied to the sum
-    transpose_op = DimShuffle(ndim=2, new_order=(1, 0))
+    transpose_op = DimShuffle(input_ndim=2, new_order=(1, 0))
     check_max_log_sum_exp(x, axis=2, dimshuffle_op=transpose_op)
 
     # If the sum is performed with keepdims=True

tests/tensor/test_elemwise.py

@@ -109,12 +109,11 @@ def test_infer_shape(self):
             ((1,), ("x", "x")),
         ]:
             i_shape = [entry if entry == 1 else None for entry in xsh]
-            ib = [(entry == 1) for entry in xsh]
             adtens = self.type(self.dtype, shape=i_shape)("x")
             adtens_val = np.ones(xsh, dtype=self.dtype)
             self._compile_and_check(
                 [adtens],
-                [self.op(ib, shuffle)(adtens)],
+                [self.op(input_ndim=len(xsh), new_order=shuffle)(adtens)],
                 [adtens_val],
                 self.op,
                 warn=False,
@@ -188,7 +187,7 @@ def test_static_shape(self):
     def test_valid_input_ndim(self):
         assert DimShuffle(input_ndim=2, new_order=(1, 0)).input_ndim == 2
 
-        with pytest.raises(TypeError, match="input_ndim must an integer"):
+        with pytest.raises(TypeError, match="input_ndim must be an integer"):
             DimShuffle(input_ndim=(True, False), new_order=(1, 0))
 
 

tests/tensor/test_extra_ops.py

@@ -481,12 +481,9 @@ def test_invalid_input(self):
         assert f([0]) == 0
 
         # Test that we cannot squeeze dimensions whose length is greater than 1
-        error_txt_1 = re.escape("SpecifyShape: Got shape (3,), expected (1,).")
-        error_txt_2 = re.escape("SpecifyShape: dim 0 of input has shape 3, expected 1")
-        match = error_txt_1 if pytensor.config.mode == "FAST_COMPILE" else error_txt_2
         with pytest.raises(
-            AssertionError,
-            match=match,
+            ValueError,
+            match="cannot reshape array of size 3 into shape ()",
         ):
             f([0, 1, 2])