Replace Shape_i with any static shape value and not just 1

Author: brandonwillard

aesara/tensor/rewriting/shape.py

@@ -1020,8 +1020,8 @@ def local_Shape_of_SpecifyShape(fgraph, node):
 @register_useless
 @register_canonicalize
 @node_rewriter([Shape_i])
-def local_Shape_i_of_broadcastable(fgraph, node):
-    """Replace ``shape_i(x, i)`` with ``1`` when ``x.broadcastable[i]`` is ``True``."""
+def local_Shape_i_ground(fgraph, node):
+    """Replace ``shape_i(x, i)`` with ``s`` when ``x.type.shape[i] == s``."""
 
     if not isinstance(node.op, Shape_i):
         return False
@@ -1031,8 +1031,9 @@ def local_Shape_i_of_broadcastable(fgraph, node):
     if not isinstance(shape_arg.type, TensorType):
         return False
 
-    if shape_arg.broadcastable[node.op.i]:
-        return [as_tensor_variable(1, dtype=np.int64)]
+    s_val = shape_arg.type.shape[node.op.i]
+    if s_val is not None:
+        return [as_tensor_variable(s_val, dtype=np.int64)]
 
 
 @register_specialize

tests/tensor/rewriting/test_shape.py

@@ -493,15 +493,15 @@ def test_local_Shape_of_SpecifyShape_partial(s1):
     assert not any(isinstance(apply.op, SpecifyShape) for apply in fgraph.apply_nodes)
 
 
-def test_local_Shape_i_of_broadcastable():
-    x = tensor(np.float64, shape=(None, 1))
+def test_local_Shape_i_ground():
+    x = tensor(np.float64, shape=(None, 2))
     s = Shape_i(1)(x)
 
     fgraph = FunctionGraph(outputs=[s], clone=False)
     _ = rewrite_graph(fgraph, clone=False)
 
     assert x not in fgraph.variables
-    assert fgraph.outputs[0].data == 1
+    assert fgraph.outputs[0].data == 2
 
     # A test for a non-`TensorType`
     class MyType(Type):