Eager optimization for no-op flatten

Author: ricardoV94

pytensor/tensor/basic.py

@@ -3081,6 +3081,10 @@ def flatten(x, ndim=1):
     else:
         dims = (-1,)
 
+    if len(dims) == _x.ndim:
+        # Nothing to ravel
+        return _x
+
     x_reshaped = _x.reshape(dims)
     shape_kept_dims = _x.type.shape[: ndim - 1]
     bcast_new_dim = builtins.all(s == 1 for s in _x.type.shape[ndim - 1 :])

tests/tensor/test_basic.py

@@ -3867,37 +3867,24 @@ class TestInferShape(utt.InferShapeTester):
     def test_Flatten(self):
         atens3 = tensor3()
         atens3_val = random(4, 5, 3)
-        for ndim in (3, 2, 1):
+        for ndim in (2, 1):
             self._compile_and_check(
                 [atens3],
                 [flatten(atens3, ndim)],
                 [atens3_val],
                 Reshape,
-                excluding=["local_useless_reshape"],
             )
 
         amat = matrix()
         amat_val = random(4, 5)
-        for ndim in (2, 1):
+        for ndim in (1,):
             self._compile_and_check(
                 [amat],
                 [flatten(amat, ndim)],
                 [amat_val],
                 Reshape,
-                excluding=["local_useless_reshape"],
             )
 
-        avec = vector()
-        avec_val = random(4)
-        ndim = 1
-        self._compile_and_check(
-            [avec],
-            [flatten(avec, ndim)],
-            [avec_val],
-            Reshape,
-            excluding=["local_useless_reshape"],
-        )
-
     def test_Eye(self):
         aiscal = iscalar()
         biscal = iscalar()