Fix bug in infer_static_shape of graphs involving the shape of scalars

Author: ricardoV94

pytensor/tensor/basic.py

@@ -1406,11 +1406,8 @@ def infer_static_shape(
     `shape` will be validated and constant folded.  As a result, this function
     can be expensive and shouldn't be used unless absolutely necessary.
 
-    It mostly exists as a hold-over from pre-static shape times, when it was
-    required in order to produce correct broadcastable arrays and prevent
-    some graphs from being unusable.  Now, it is no longer strictly required,
-    so don't use it unless you want the same shape graphs to be rewritten
-    multiple times during graph construction.
+    It is often needed for `Op`s whose static shape and broadcastable flags
+    depend on the values of their inputs, such as `Alloc` and `RandomVariable`.
 
     Returns
     -------

pytensor/tensor/rewriting/shape.py

@@ -992,12 +992,17 @@ def local_merge_consecutive_specify_shape(fgraph, node):
     return [specify_shape(inner_obj, shape)]
 
 
+_empty_shape = constant([], dtype="int64")
+
+
 @register_infer_shape
 @node_rewriter([Shape])
 def local_shape_ground(fgraph, node):
     """Rewrite shape(x) -> make_vector(x.type.shape) when this is constant."""
     [x] = node.inputs
     static_shape = x.type.shape
+    if len(static_shape) == 0:
+        return [_empty_shape]
     if not any(dim is None for dim in static_shape):
         return [stack([constant(dim, dtype="int64") for dim in static_shape])]
 

tests/tensor/test_basic.py

@@ -908,7 +908,7 @@ def test_runtime_broadcast(self, mode):
         self.check_runtime_broadcast(mode)
 
 
-def test_infer_shape():
+def test_infer_static_shape():
     with pytest.raises(TypeError, match="^Shapes must be scalar integers.*"):
         infer_static_shape([constant(1.0)])
 
@@ -925,6 +925,10 @@ def test_infer_shape():
     sh, static_shape = infer_static_shape(specify_size)
     assert static_shape == (1,)
 
+    x = scalar("x")
+    sh, static_shape = infer_static_shape([x.size])
+    assert static_shape == (1,)
+
 
 # This is slow for the ('int8', 3) version.
 def test_eye():