Create scalarize rewrite pass

Author: aseyboldt

pytensor/compile/mode.py

@@ -250,6 +250,11 @@ def apply(self, fgraph):
 # misc special cases for speed that break canonicalization
 optdb.register("uncanonicalize", EquilibriumDB(), "fast_run", position=3)
 
+# Turn tensor operations to scalar operations where possible.
+# This is currently marked as numba-only, but this could be changed
+# in the future.
+optdb.register("scalarize", EquilibriumDB(), "numba_only", position=3.1)
+
 # misc special cases for speed that are dependent on the device.
 optdb.register(
     "specialize_device", EquilibriumDB(), "fast_compile", "fast_run", position=48.6
@@ -459,20 +464,42 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
 # FunctionMaker, the Mode will be taken from this dictionary using the
 # string as the key
 # Use VM_linker to allow lazy evaluation by default.
-FAST_COMPILE = Mode(VMLinker(use_cloop=False, c_thunks=False), "fast_compile")
+FAST_COMPILE = Mode(
+    VMLinker(use_cloop=False, c_thunks=False),
+    RewriteDatabaseQuery(
+        include=["fast_compile"],
+        exclude=["numba_only"],
+    ),
+)
 if config.cxx:
-    FAST_RUN = Mode("cvm", "fast_run")
+    FAST_RUN = Mode(
+        "cvm",
+        RewriteDatabaseQuery(
+            include=["fast_run"],
+            exclude=["numba_only"],
+        ),
+    )
 else:
-    FAST_RUN = Mode("vm", "fast_run")
+    FAST_RUN = Mode(
+        "vm",
+        RewriteDatabaseQuery(
+            include=["fast_run"],
+            exclude=["numba_only"],
+        ),
+    )
 
 JAX = Mode(
     JAXLinker(),
-    RewriteDatabaseQuery(include=["fast_run", "jax"], exclude=["cxx_only", "BlasOpt"]),
+    RewriteDatabaseQuery(
+        include=["fast_run", "jax"],
+        exclude=["cxx_only", "BlasOpt", "numba_only"],
+    ),
 )
+
 NUMBA = Mode(
     NumbaLinker(),
     RewriteDatabaseQuery(
-        include=["fast_run", "fast_run_numba", "fast_compile_numba"],
+        include=["fast_run", "numba_only"],
         exclude=["cxx_only", "BlasOpt"],
     ),
 )

pytensor/tensor/rewriting/basic.py

@@ -186,6 +186,23 @@ def register(inner_rewriter: Union[RewriteDatabase, Rewriter]):
         return node_rewriter
 
 
+def register_scalarize(
+    node_rewriter: Union[RewriteDatabase, NodeRewriter, str], *tags: str, **kwargs
+):
+    if isinstance(node_rewriter, str):
+
+        def register(inner_rewriter: Union[RewriteDatabase, Rewriter]):
+            return register_specialize(inner_rewriter, node_rewriter, *tags, **kwargs)
+
+        return register
+    else:
+        name = kwargs.pop("name", None) or node_rewriter.__name__
+        compile.optdb["scalarize"].register(
+            name, node_rewriter, "fast_run", "fast_compile", *tags, **kwargs
+        )
+        return node_rewriter
+
+
 def register_uncanonicalize(
     node_rewriter: Union[RewriteDatabase, NodeRewriter, str], *tags: str, **kwargs
 ):
@@ -226,30 +243,36 @@ def register(inner_rewriter: Union[RewriteDatabase, Rewriter]):
 
 @register_canonicalize
 @register_specialize
+@register_scalarize
 @node_rewriter([TensorFromScalar])
 def local_tensor_scalar_tensor(fgraph, node):
     """tensor_from_scalar(scalar_from_tensor(x)) -> x"""
-    if isinstance(node.op, TensorFromScalar):
-        s = node.inputs[0]
-        if s.owner and isinstance(s.owner.op, ScalarFromTensor):
-            t = s.owner.inputs[0]
+    s = node.inputs[0]
+    if s.owner and isinstance(s.owner.op, ScalarFromTensor):
+        t = s.owner.inputs[0]
 
-            # We don't need to copy over any stack traces here
-            return [t]
+        # We don't need to copy over any stack traces here
+        return [t]
 
 
 @register_canonicalize
 @register_specialize
+@register_scalarize
 @node_rewriter([ScalarFromTensor])
 def local_scalar_tensor_scalar(fgraph, node):
-    """scalar_from_tensor(tensor_from_scalar(x)) -> x"""
-    if isinstance(node.op, ScalarFromTensor):
-        t = node.inputs[0]
-        if t.owner and isinstance(t.owner.op, TensorFromScalar):
-            s = t.owner.inputs[0]
-
-            # We don't need to copy over any stack traces here
-            return [s]
+    """scalar_from_tensor(tensor_from_scalar(x)) -> x
+
+    and scalar_from_tensor(TensorConstant(x)) -> x
+    """
+    t = node.inputs[0]
+    if t.owner and isinstance(t.owner.op, TensorFromScalar):
+        s = t.owner.inputs[0]
+
+        # We don't need to copy over any stack traces here
+        return [s]
+    if isinstance(t, TensorConstant):
+        assert t.ndim == 0
+        return [aes.constant(t.value.item(), t.name, t.dtype)]
 
 
 @register_specialize("local_alloc_elemwise")

pytensor/tensor/rewriting/elemwise.py

@@ -381,12 +381,8 @@ def is_dimshuffle_useless(new_order, input):
 
 
 @node_rewriter([Elemwise])
-def local_elemwise_lift_scalars(fgraph, node):
+def elemwise_to_scalar(fgraph, node):
     op = node.op
-
-    if not isinstance(op, Elemwise):
-        return False
-
     if not all(input.ndim == 0 for input in node.inputs):
         return False
 
@@ -397,11 +393,12 @@ def local_elemwise_lift_scalars(fgraph, node):
     return [as_tensor_variable(out) for out in op.scalar_op.make_node(*scalars).outputs]
 
 
-compile.optdb["specialize"].register(
-    "local_elemwise_lift_scalars",
-    local_elemwise_lift_scalars,
-    "fast_run_numba",
-    "fast_compile_numba",
+compile.optdb["scalarize"].register(
+    "local_elemwise_to_scalar",
+    elemwise_to_scalar,
+    "fast_run",
+    "fast_compile",
+    "numba_only",
 )
 
 
@@ -411,9 +408,6 @@ def push_elemwise_constants(fgraph, node):
     contained scalar op.
     """
     op = node.op
-    if not isinstance(op, Elemwise):
-        return False
-
     if any(op.inplace_pattern):
         return False
 
@@ -467,8 +461,7 @@ def is_constant_scalar(x):
 compile.optdb["post_fusion"].register(
     "push_elemwise_constants",
     push_elemwise_constants,
-    "fast_run_numba",
-    "fast_compile_numba",
+    "numba_only",
 )
 
 

pytensor/tensor/rewriting/math.py

@@ -86,6 +86,7 @@
     encompasses_broadcastable,
     local_fill_sink,
     register_canonicalize,
+    register_scalarize,
     register_specialize,
     register_specialize_device,
     register_stabilize,
@@ -1568,6 +1569,18 @@ def local_op_of_op(fgraph, node):
                 return [combined(node_inps.owner.inputs[0])]
 
 
+@register_scalarize
+@node_rewriter([Sum])
+def local_sum_of_makevector(fgraph, node):
+    (array,) = node.inputs
+    if not array.owner or not isinstance(array.owner.op, MakeVector):
+        return False
+
+    values = array.owner.inputs
+    summed = aes.add(*values)
+    return [as_tensor_variable(summed)]
+
+
 ALL_REDUCE = (
     [
         CAReduce,

pytensor/tensor/rewriting/subtensor.py

@@ -469,7 +469,8 @@ def local_subtensor_lift(fgraph, node):
             return [rbcast_subt_x]
 
 
-@register_specialize
+@register_stabilize("cxx_only")
+@register_canonicalize("cxx_only")
 @node_rewriter([Subtensor])
 def local_subtensor_merge(fgraph, node):
     """

tests/link/numba/test_scan.py

@@ -10,7 +10,6 @@
 from pytensor.scan.op import Scan
 from pytensor.scan.utils import until
 from pytensor.tensor import log, vector
-from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.random.utils import RandomStream
 from tests import unittest_tools as utt
 from tests.link.numba.test_basic import compare_numba_and_py
@@ -437,8 +436,4 @@ def test_inner_graph_optimized():
         node for node in f.maker.fgraph.apply_nodes if isinstance(node.op, Scan)
     ]
     inner_scan_nodes = scan_node.op.fgraph.apply_nodes
-    assert len(inner_scan_nodes) == 1
-    (inner_scan_node,) = scan_node.op.fgraph.apply_nodes
-    assert isinstance(inner_scan_node.op, Elemwise) and isinstance(
-        inner_scan_node.op.scalar_op, Log1p
-    )
+    assert any(isinstance(node.op, Log1p) for node in inner_scan_nodes)