Fix bug in local_reshape_to_dimshuffle

Author: ricardoV94

pytensor/tensor/rewriting/shape.py

@@ -966,16 +966,15 @@ def local_reshape_to_dimshuffle(fgraph, node):
  inp, output_shape = node.inputs
  [output] = node.outputs
 
- # Remove any broadcastable dimensions from the input
- squeeze_axes = [i for i, bcast in enumerate(inp.type.broadcastable) if bcast]
-
  # Trivial case, all dimensions of input/output are known to be broadcastable:
  # there's nothing to reshape
  if all(inp.type.broadcastable) or all(output.type.broadcastable):
+ squeeze_axes = tuple(range(inp.type.ndim))
  new_output_shape = []
  expand_axes = tuple(range(output.type.ndim))
 
  else:
+ squeeze_axes = [i for i, bcast in enumerate(inp.type.broadcastable) if bcast]
  unpacked_shape = _unpack_shape_vector(output_shape)
  new_output_shape = []
  expand_axes = []

tests/tensor/rewriting/test_shape.py

@@ -445,6 +445,15 @@ def test_squeeze_of_alloc(self):
  new_out = rewrite_graph(out, include=("canonicalize", "ShapeOpt"))
  assert equal_computations([new_out], [pt.alloc(x, 12, 9)], strict_dtype=False)
 
+ def test_reshape_implies_size_1_input(self):
+ x = pt.matrix("x", shape=(None, None))
+ out = pt.reshape(x, (1, 1, 1))
+
+ new_out = rewrite_graph(out, include=("canonicalize",))
+ assert equal_computations(
+ [new_out], [x.dimshuffle("x", "x", "x")], strict_dtype=False
+ )
+
 
 def test_expand_dims_squeeze_reshape_fusion():
  x = pt.tensor("x", shape=(1, 9))