Faster infer_shape

Author: ricardoV94

pytensor/tensor/basic.py

@@ -22,9 +22,11 @@
 from pytensor import compile, config, printing
 from pytensor import scalar as aes
 from pytensor.gradient import DisconnectedType, grad_undefined
+from pytensor.graph import RewriteDatabaseQuery
 from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
+from pytensor.graph.rewriting.db import TopoDB, EquilibriumDB
 from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.graph.type import HasShape, Type
 from pytensor.link.c.op import COp
@@ -1356,6 +1358,25 @@ def identity_like(x, dtype: Optional[Union[str, np.generic, np.dtype]] = None):
     return eye(_x.shape[0], _x.shape[1], k=0, dtype=dtype)
 
 
+infer_shape_db = EquilibriumDB()
+
+
+def register_infer_shape(rewrite, *tags, **kwargs):
+    if isinstance(rewrite, str):
+
+        def register(inner_lopt):
+            return register_infer_shape(inner_lopt, rewrite, *tags, **kwargs)
+
+        return register
+    else:
+        name = kwargs.pop("name", None) or rewrite.__name__
+
+        infer_shape_db.register(
+            name, rewrite, *tags, "infer_shape", **kwargs
+        )
+        return rewrite
+
+
 def infer_static_shape(
     shape: Union[Variable, Sequence[Union[Variable, int]]]
 ) -> tuple[Sequence["TensorLike"], Sequence[Optional[int]]]:
@@ -1391,13 +1412,13 @@ def check_type(s):
         raise TypeError(f"Shapes must be scalar integers; got {s_as_str}")
 
     sh = [check_type(as_tensor_variable(s, ndim=0)) for s in shape]
-
-    shape_fg = FunctionGraph(
-        outputs=sh,
-        features=[ShapeFeature()],
-        clone=True,
-    )
-    folded_shape = rewrite_graph(shape_fg, custom_rewrite=topo_constant_folding).outputs
+    shape_fg = FunctionGraph(outputs=sh, features=[ShapeFeature()], clone=True)
+    with config.change_flags(optdb__max_use_ratio=10, cxx=""):
+        query_rewrites = infer_shape_db.query(RewriteDatabaseQuery(include=("infer_shape",)))
+        query_rewrites.rewrite(shape_fg)
+        if not all(isinstance(s, Constant) for s in shape_fg.outputs):
+            topo_constant_folding.rewrite(shape_fg)
+    folded_shape = shape_fg.outputs
     static_shape = tuple(
         s.data.item() if isinstance(s, Constant) else None for s in folded_shape
     )

pytensor/tensor/rewriting/basic.py

@@ -61,7 +61,7 @@
     switch,
     tensor_copy,
     zeros,
-    zeros_like,
+    zeros_like, register_infer_shape,
 )
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
@@ -420,6 +420,7 @@ def local_fill_to_alloc(fgraph, node):
 )
 
 
+@register_infer_shape
 @register_canonicalize("fast_compile", "shape_unsafe")
 @register_useless("shape_unsafe")
 @node_rewriter([fill])
@@ -441,6 +442,7 @@ def local_useless_fill(fgraph, node):
         return [v]
 
 
+@register_infer_shape
 @register_specialize("shape_unsafe")
 @register_stabilize("shape_unsafe")
 @register_canonicalize("shape_unsafe")
@@ -530,6 +532,7 @@ def local_alloc_empty_to_zeros(fgraph, node):
 )
 
 
+@register_infer_shape
 @register_useless
 @register_canonicalize("fast_compile")
 @register_specialize
@@ -806,6 +809,7 @@ def local_remove_all_assert(fgraph, node):
 )
 
 
+@register_infer_shape
 @register_specialize
 @register_canonicalize
 @register_useless
@@ -826,6 +830,7 @@ def local_join_1(fgraph, node):
 
 
 # TODO: merge in local_useless_join
+@register_infer_shape
 @register_useless
 @register_specialize
 @register_canonicalize
@@ -1066,6 +1071,7 @@ def local_merge_switch_same_cond(fgraph, node):
     ]
 
 
+@register_infer_shape
 @register_useless
 @register_canonicalize
 @register_specialize
@@ -1149,6 +1155,7 @@ def constant_folding(fgraph, node):
 register_specialize(topo_constant_folding, "fast_compile", final_rewriter=True)
 
 
+@register_infer_shape
 @register_canonicalize("fast_compile")
 @register_useless("fast_compile")
 @node_rewriter(None)
@@ -1157,6 +1164,7 @@ def local_view_op(fgraph, node):
         return node.inputs
 
 
+@register_infer_shape
 @register_useless
 @register_canonicalize
 @register_stabilize

pytensor/tensor/rewriting/math.py

@@ -33,7 +33,7 @@
     get_underlying_scalar_constant_value,
     ones_like,
     switch,
-    zeros_like,
+    zeros_like, register_infer_shape,
 )
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
@@ -1745,6 +1745,7 @@ def local_reduce_join(fgraph, node):
         return [ret]
 
 
+@register_infer_shape
 @register_canonicalize("fast_compile", "local_cut_useless_reduce")
 @register_useless("local_cut_useless_reduce")
 @node_rewriter(ALL_REDUCE)

pytensor/tensor/rewriting/shape.py

@@ -17,6 +17,7 @@
     copy_stack_trace,
     node_rewriter,
 )
+from pytensor.graph.rewriting.db import TopoDB
 from pytensor.graph.utils import InconsistencyError, get_variable_trace_string
 from pytensor.tensor.basic import (
     MakeVector,
@@ -25,7 +26,7 @@
     constant,
     extract_constant,
     get_underlying_scalar_constant_value,
-    stack,
+    stack, register_infer_shape,
 )
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError, ShapeError
@@ -964,6 +965,7 @@ def local_reshape_lift(fgraph, node):
         return [e]
 
 
+@register_infer_shape
 @register_useless
 @register_canonicalize
 @node_rewriter([SpecifyShape])
@@ -990,6 +992,25 @@ def local_merge_consecutive_specify_shape(fgraph, node):
     return [specify_shape(inner_obj, shape)]
 
 
+@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 not any(dim is None for dim in static_shape):
+        return [stack([constant(dim, dtype="int64") for dim in static_shape])]
+
+
+# Don't register in FAST_RUN, as we don't want to get rid of the shape before putting in Shape_i
+pytensor.compile.mode.optdb["useless"].register(
+    "local_shape_ground",
+    local_shape_ground,
+    position="last",
+)
+
+
+@register_infer_shape
 @register_useless
 @register_canonicalize
 @node_rewriter([Shape])
@@ -1014,6 +1035,7 @@ def local_Shape_of_SpecifyShape(fgraph, node):
     return [stack(shape).astype(np.int64)]
 
 
+@register_infer_shape
 @register_canonicalize
 @register_specialize
 @node_rewriter([SpecifyShape])
@@ -1060,6 +1082,7 @@ def local_specify_shape_lift(fgraph, node):
         return new_out
 
 
+@register_infer_shape
 @register_useless
 @register_canonicalize
 @node_rewriter([Shape_i])
@@ -1079,6 +1102,7 @@ def local_Shape_i_ground(fgraph, node):
         return [as_tensor_variable(s_val, dtype=np.int64)]
 
 
+@register_infer_shape
 @register_specialize
 @register_canonicalize
 @node_rewriter([Shape])

pytensor/tensor/rewriting/subtensor.py

@@ -26,7 +26,7 @@
     concatenate,
     extract_constant,
     get_underlying_scalar_constant_value,
-    switch,
+    switch, register_infer_shape,
 )
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
@@ -328,6 +328,7 @@ def local_subtensor_of_dot(fgraph, node):
     return [r]
 
 
+@register_infer_shape
 @register_useless
 @register_canonicalize
 @register_specialize
@@ -599,6 +600,7 @@ def local_subtensor_remove_broadcastable_index(fgraph, node):
         return [node.inputs[0].dimshuffle(tuple(remain_dim))]
 
 
+@register_infer_shape
 @register_useless
 @register_canonicalize
 @register_specialize
@@ -707,6 +709,7 @@ def local_subtensor_inc_subtensor(fgraph, node):
             return
 
 
+@register_infer_shape
 @register_specialize
 @register_canonicalize("fast_compile")
 @register_useless
@@ -785,6 +788,7 @@ def local_subtensor_make_vector(fgraph, node):
             pass
 
 
+@register_infer_shape
 @register_useless
 @register_canonicalize
 @register_specialize
@@ -1461,6 +1465,7 @@ def local_adv_sub1_adv_inc_sub1(fgraph, node):
     return [r2]
 
 
+@register_infer_shape
 @register_specialize
 @register_stabilize
 @register_canonicalize

tests/tensor/test_blockwise.py

@@ -358,8 +358,10 @@ def test_batched_mvnormal_logp_and_dlogp(mu_batch_shape, cov_batch_shape, benchm
     logp = norm - 0.5 * quaddist - logdet
     dlogp = grad(logp.sum(), wrt=[value, mu, cov])
 
+    # pytensor.dprint([logp, *dlogp], print_type=True)
     fn = pytensor.function([value, mu, cov], [logp, *dlogp])
-    benchmark(fn, *test_values)
+    pytensor.dprint(fn, print_type=True)
+    # benchmark(fn, *test_values)
 
 
 def test_cop_with_params():