Create new rewrites for elemwise

There is no need for an Elemwise Op if all inputs have rank 0.
And we don't need to use scalar constants as inputs of the
Elemwise, they can be inputs for the scalar_op.

Author: aseyboldt

pytensor/_version.py

@@ -92,7 +92,7 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env=
                 env=env,
                 stdout=subprocess.PIPE,
                 stderr=(subprocess.PIPE if hide_stderr else None),
-                **popen_kwargs
+                **popen_kwargs,
             )
             break
         except OSError:

pytensor/compile/mode.py

@@ -255,6 +255,16 @@ def apply(self, fgraph):
     "specialize_device", EquilibriumDB(), "fast_compile", "fast_run", position=48.6
 )  # must be after gpu stuff at 48.5
 
+# Must be before add_destroy_handler
+optdb.register(
+    "elemwise_fusion",
+    SequenceDB(),
+    "fast_run",
+    "fusion",
+    "local_elemwise_fusion",
+    position=49,
+)
+
 # especially constant merge
 optdb.register("merge2", MergeOptimizer(), "fast_run", "merge", position=49)
 
@@ -453,7 +463,10 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
 )
 NUMBA = Mode(
     NumbaLinker(),
-    RewriteDatabaseQuery(include=["fast_run"], exclude=["cxx_only", "BlasOpt"]),
+    RewriteDatabaseQuery(
+        include=["fast_run", "fast_run_numba", "fast_compile_numba"],
+        exclude=["cxx_only", "BlasOpt"],
+    ),
 )
 
 

pytensor/tensor/rewriting/elemwise.py

@@ -18,8 +18,8 @@
     in2out,
     node_rewriter,
 )
-from pytensor.graph.rewriting.db import SequenceDB
 from pytensor.graph.utils import InconsistencyError, MethodNotDefined, TestValueError
+from pytensor.tensor import as_tensor_variable
 from pytensor.tensor.basic import MakeVector, alloc, cast, get_scalar_constant_value
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
@@ -380,6 +380,99 @@ def is_dimshuffle_useless(new_order, input):
     return is_useless
 
 
+@node_rewriter([Elemwise])
+def local_elemwise_lift_scalars(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
+
+    scalars = [aes.as_scalar(input) for input in node.inputs]
+
+    # TODO Something like
+    # copy_stack_trace(node.outputs[0], new_res)
+    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",
+)
+
+
+@node_rewriter([Elemwise])
+def push_elemwise_constants(fgraph, node):
+    """Push constant scalars from inputs to elemwise to inputs of the
+    contained scalar op.
+    """
+    op = node.op
+
+    if not isinstance(op, Elemwise):
+        return False
+
+    if any(op.inplace_pattern):
+        return False
+
+    if not isinstance(node.op.scalar_op, aes.Composite):
+        return False
+
+    def is_constant_scalar(x):
+        return isinstance(x, TensorConstant) and all(x.broadcastable)
+
+    push_idxs = []
+    push_values = []
+    keep_values = []
+    for i, input in enumerate(node.inputs):
+        if is_constant_scalar(input):
+            push_idxs.append(i)
+            val = input.value
+            push_values.append(aes.constant(val.item(), dtype=val.dtype))
+        elif (
+            input.owner
+            and isinstance(input.owner.op, DimShuffle)
+            and is_constant_scalar(input.owner.inputs[0])
+        ):
+            push_idxs.append(i)
+            val = input.owner.inputs[0].value
+            push_values.append(aes.constant(val.item(), dtype=val.dtype))
+        else:
+            keep_values.append(input)
+
+    if not push_values:
+        return False
+
+    inner_graph = node.op.scalar_op.fgraph
+    to_replace = [input for i, input in enumerate(inner_graph.inputs) if i in push_idxs]
+
+    # Clone the inner graph, it might be used somewhere else
+    inner_graph, mapping = inner_graph.clone_get_equiv()
+    inner_graph.replace_all(
+        (mapping[old], new) for old, new in zip(to_replace, push_values)
+    )
+
+    new_inputs = [
+        input for i, input in enumerate(inner_graph.inputs) if i not in push_idxs
+    ]
+    return (
+        Elemwise(scalar_op=aes.Composite(new_inputs, inner_graph.outputs))
+        .make_node(*keep_values)
+        .outputs
+    )
+
+
+compile.optdb["specialize"].register(
+    "push_elemwise_constants",
+    push_elemwise_constants,
+    "fast_run_numba",
+    "fast_compile_numba",
+)
+
+
 @register_canonicalize
 @register_specialize
 @node_rewriter([DimShuffle])
@@ -898,34 +991,13 @@ def print_profile(cls, stream, prof, level=0):
         print(blanc, " time_toposort", prof[7], file=stream)
 
 
-if config.tensor__local_elemwise_fusion:
-    # Must be after gpu(48.5) and before AddDestroyHandler(49.5)
-    fuse_seqopt = SequenceDB()
-    fuse_seqopt.register(
-        "composite_elemwise_fusion",
-        FusionOptimizer(local_elemwise_fusion),
-        "fast_run",
-        "fusion",
-        position=1,
-    )
-    compile.optdb.register(  # type: ignore
-        "elemwise_fusion",
-        fuse_seqopt,
-        "fast_run",
-        "fusion",
-        "local_elemwise_fusion",
-        "FusionOptimizer",
-        position=49,
-    )
-else:
-    compile.optdb.register(  # type: ignore
-        "elemwise_fusion",
-        FusionOptimizer(local_elemwise_fusion),
-        "fusion",
-        "local_elemwise_fusion",
-        "FusionOptimizer",
-        position=49,
-    )
+compile.optdb["elemwise_fusion"].register(  # type: ignore
+    "composite_elemwise_fusion",
+    FusionOptimizer(local_elemwise_fusion),
+    "fast_run",
+    "fusion",
+    position=1,
+)
 
 
 @register_canonicalize

pytensor/tensor/rewriting/math.py

@@ -8,6 +8,7 @@
 
 import pytensor.scalar.basic as aes
 import pytensor.scalar.math as aes_math
+from pytensor import compile
 from pytensor.graph.basic import Constant, Variable
 from pytensor.graph.rewriting.basic import (
     NodeRewriter,
@@ -91,7 +92,7 @@
     register_uncanonicalize,
     register_useless,
 )
-from pytensor.tensor.rewriting.elemwise import FusionOptimizer, fuse_seqopt
+from pytensor.tensor.rewriting.elemwise import FusionOptimizer
 from pytensor.tensor.shape import Shape, Shape_i
 from pytensor.tensor.subtensor import Subtensor
 from pytensor.tensor.type import (
@@ -2922,7 +2923,7 @@ def local_add_mul_fusion(fgraph, node):
         return [output]
 
 
-fuse_seqopt.register(
+compile.optdb["elemwise_fusion"].register(
     "local_add_mul_fusion",
     FusionOptimizer(local_add_mul_fusion),
     "fast_run",

pytensor/tensor/rewriting/subtensor.py

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