Avoid no-op DimShuffle

Author: ricardoV94

pytensor/tensor/variable.py

@@ -349,6 +349,9 @@ def dimshuffle(self, *pattern):
         if (len(pattern) == 1) and (isinstance(pattern[0], list | tuple | np.ndarray)):
             pattern = pattern[0]
         ds_op = pt.elemwise.DimShuffle(input_ndim=self.type.ndim, new_order=pattern)
+        if ds_op.new_order == tuple(range(self.type.ndim)):
+            # No-op
+            return self
         return ds_op(self)
 
     def flatten(self, ndim=1):

tests/tensor/random/rewriting/test_basic.py

@@ -950,7 +950,7 @@ def test_Dimshuffle_lift_restrictions():
             1e-7,
         ),
         (
-            (0, 1, 2),
+            (0, 2, 1),
             True,
             normal,
             (np.array(0).astype(config.floatX), np.array(1e-6).astype(config.floatX)),

tests/tensor/rewriting/test_elemwise.py

@@ -148,7 +148,7 @@ def test_recursive_lift(self):
 
     def test_useless_dimshuffle(self):
         x, *_ = inputs()
-        e = ds(x, (0, 1))
+        e = DimShuffle(new_order=(0, 1), input_ndim=2)(x)
         g = FunctionGraph([x], [e], clone=False)
         assert isinstance(g.outputs[0].owner.op, DimShuffle)
         dimshuffle_lift.rewrite(g)