Rewrite batched dots that do not reduce as multiplication

pytensor/graph/rewriting/__init__.py

@@ -0,0 +1,4 @@
+from pytensor.graph.rewriting.utils import rewrite_graph
+
+
+all = ("rewrite_graph",)

pytensor/tensor/math.py

@@ -29,7 +29,7 @@
     stack,
     switch,
 )
-from pytensor.tensor.blockwise import Blockwise, vectorize_node_fallback
+from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import (
     CAReduce,
     Elemwise,
@@ -2726,6 +2726,22 @@
     return log(sum(exp(x), axis=axis, keepdims=keepdims))
 
 
+# Predefine all batched variations of Dot
+_inner_prod = Blockwise(
+    _dot,
+    signature="(n),(n)->()",
+)
+
+_matrix_vec_prod = Blockwise(
+    _dot,
+    signature="(m,k),(k)->(m)",
+)
+
+_vec_matrix_prod = Blockwise(
+    _dot,
+    signature="(k),(k,n)->(n)",
+)
+
 _matrix_matrix_matmul = Blockwise(
     _dot,
     signature="(m,k),(k,n)->(m,n)",
@@ -2795,14 +2811,24 @@
 
 
 @_vectorize_node.register(Dot)
-def vectorize_node_dot_to_matmul(op, node, batched_x, batched_y):
+def vectorize_node_dot(op, node, batched_x, batched_y):
     old_x, old_y = node.inputs
-    if old_x.type.ndim == 2 and old_y.type.ndim == 2:
-        # If original input is equivalent to a matrix-matrix product,
-        # return specialized Matmul Op to avoid unnecessary new Ops.
-        return matmul(batched_x, batched_y).owner
-    else:
-        return vectorize_node_fallback(op, node, batched_x, batched_y)
+    old_x_ndim = old_x.type.ndim
+    old_y_ndim = old_y.type.ndim
+    match (old_x_ndim, old_y_ndim):
+        case (1, 1):
+            batch_op = _inner_prod
+        case (2, 1):
+            batch_op = _matrix_vec_prod
+        case (1, 2):
+            batch_op = _vec_matrix_prod
+        case (2, 2):
+            batch_op = _matrix_matrix_matmul
+        case _:
+            raise ValueError(
+                f"Core dot Op should have 1D or 2D inputs, got {old_x_ndim}D and {old_y_ndim}D."
+            )
+    return batch_op(batched_x, batched_y).owner
 
 
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):

pytensor/tensor/rewriting/math.py

@@ -44,6 +44,10 @@
     Prod,
     Sum,
     _conj,
+    _inner_prod,
+    _matrix_matrix_matmul,
+    _matrix_vec_prod,
+    _vec_matrix_prod,
     add,
     digamma,
     dot,
@@ -242,6 +246,62 @@
     return None
 
 
+@register_canonicalize
+@register_specialize
+@node_rewriter([_inner_prod, _matrix_vec_prod, _vec_matrix_prod, _matrix_matrix_matmul])
+def local_blockwise_dot_to_mul(fgraph, node):
+    """Rewrite blockwise dots that correspond to multiplication without summation.
+
+    We don't touch the regular dot, to not interfere with the BLAS optimizations.
+    """
+    a, b = node.inputs
+    a_static_shape = a.type.shape
+    b_static_shape = b.type.shape
+    core_a_ndim = len(node.op.inputs_sig[0])
+    core_b_ndim = len(node.op.inputs_sig[1])
+
+    if core_a_ndim > 2 or core_b_ndim > 2:
+        # Shouldn't happen, but here just in case
+        return None
+
+    if core_b_ndim == 1:
+        if a_static_shape[-1] == 1 or b_static_shape[-1] == 1:
+            if core_a_ndim == 1:
+                # inner product: (..., 1) * (..., 1) -> (...)
+                # just squeeze the last dimensions of a and b
+                new_a = a.squeeze(-1)
+                new_b = b.squeeze(-1)
+            else:
+                # matrix vector product: (..., m, 1) * (..., 1) -> (..., m)
+                # the last dimension of b is already aligned for the elemwise multiplication
+                # after we squeeze the last dimension of a
+                new_a = a.squeeze(-1)
+                new_b = b
+        else:
+            return None
+
+    else:
+        if a_static_shape[-1] == 1 or b_static_shape[-2] == 1:
+            if core_a_ndim == 1:
+                # vector_matrix product: (..., 1) * (..., 1, n) -> (..., n)
+                # the last dimension of a is already aligned for the elemwise multiplication
+                # after we squeeze the one to last dimension of b
+                new_a = a
+                new_b = b.squeeze(-2)
+            else:
+                # matrix matrix product: (..., m, 1) * (..., 1, n) -> (..., m, n)
+                # the dimensions of a and b are already aligned for the elemwise multiplication
+                new_a = a
+                new_b = b
+        else:
+            return None
+
+    new_a = copy_stack_trace(a, new_a)
+    new_b = copy_stack_trace(b, new_b)
+    new_out = copy_stack_trace(node.out, mul(new_a, new_b))
+    return [new_out]
+
+
 def is_inverse_pair(node_op, prev_op, inv_pair):
     """
     Given two consecutive operations, check if they are the

tests/link/jax/test_elemwise.py

@@ -15,11 +15,11 @@
 from pytensor.tensor.special import SoftmaxGrad, log_softmax, softmax
 from pytensor.tensor.type import matrix, tensor, vector, vectors
 from tests.link.jax.test_basic import compare_jax_and_py
-from tests.tensor.test_elemwise import TestElemwise
+from tests.tensor.test_elemwise import check_elemwise_runtime_broadcast
 
 
 def test_elemwise_runtime_broadcast():
-    TestElemwise.check_runtime_broadcast(get_mode("JAX"))
+    check_elemwise_runtime_broadcast(get_mode("JAX"))
 
 
 def test_jax_Dimshuffle():

tests/link/jax/test_tensor_basic.py

@@ -14,7 +14,7 @@
 from pytensor.graph.op import get_test_value
 from pytensor.tensor.type import iscalar, matrix, scalar, vector
 from tests.link.jax.test_basic import compare_jax_and_py
-from tests.tensor.test_basic import TestAlloc
+from tests.tensor.test_basic import check_alloc_runtime_broadcast
 
 
 def test_jax_Alloc():
@@ -54,7 +54,7 @@ def compare_shape_dtype(x, y):
 
 
 def test_alloc_runtime_broadcast():
-    TestAlloc.check_runtime_broadcast(get_mode("JAX"))
+    check_alloc_runtime_broadcast(get_mode("JAX"))
 
 
 def test_jax_MakeVector():

tests/link/numba/test_elemwise.py

@@ -24,7 +24,10 @@
     scalar_my_multi_out,
     set_test_value,
 )
-from tests.tensor.test_elemwise import TestElemwise, careduce_benchmark_tester
+from tests.tensor.test_elemwise import (
+    careduce_benchmark_tester,
+    check_elemwise_runtime_broadcast,
+)
 
 
 rng = np.random.default_rng(42849)
@@ -124,7 +127,7 @@ def test_Elemwise(inputs, input_vals, output_fn, exc):
 
 @pytest.mark.xfail(reason="Logic had to be reversed due to surprising segfaults")
 def test_elemwise_runtime_broadcast():
-    TestElemwise.check_runtime_broadcast(get_mode("NUMBA"))
+    check_elemwise_runtime_broadcast(get_mode("NUMBA"))
 
 
 def test_elemwise_speed(benchmark):

tests/link/numba/test_tensor_basic.py

@@ -16,7 +16,7 @@
     compare_shape_dtype,
     set_test_value,
 )
-from tests.tensor.test_basic import TestAlloc
+from tests.tensor.test_basic import check_alloc_runtime_broadcast
 
 
 pytest.importorskip("numba")
@@ -52,7 +52,7 @@ def test_Alloc(v, shape):
 
 
 def test_alloc_runtime_broadcast():
-    TestAlloc.check_runtime_broadcast(get_mode("NUMBA"))
+    check_alloc_runtime_broadcast(get_mode("NUMBA"))
 
 
 def test_AllocEmpty():

tests/tensor/rewriting/test_math.py

@@ -16,7 +16,8 @@
 from pytensor.compile.mode import Mode, get_default_mode, get_mode
 from pytensor.compile.ops import DeepCopyOp, deep_copy_op
 from pytensor.configdefaults import config
-from pytensor.graph.basic import Apply, equal_computations
+from pytensor.graph import vectorize_graph
+from pytensor.graph.basic import Apply, ancestors, equal_computations
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import (
     SequentialNodeRewriter,
@@ -4590,3 +4591,53 @@ def test_pow_1_rewrite(shape):
 
     x_val = np.random.random(shape).astype(config.floatX)
     np.testing.assert_allclose(z.eval({x: x_val}), f(x_val))
+
+
+@pytest.mark.parametrize(
+    "a_shape,b_shape",
+    [
+        ((1,), (1,)),
+        ((3, 1), (1,)),
+        ((1,), (1, 3)),
+        ((3, 1), (1, 3)),
+    ],
+    ids=str,
+)
+@pytest.mark.parametrize("batched", (False, True))
+def test_local_dot_to_mul(batched, a_shape, b_shape):
+    a = tensor("a", shape=a_shape)
+    b = tensor("b", shape=b_shape)
+
+    out = dot(a, b)
+    if batched:
+        batch_a = tensor("batch_a", shape=(1, 5, *a_shape))
+        batch_b = tensor("batch_b", shape=(7, 1, *b_shape))
+        out = vectorize_graph(out, {a: batch_a, b: batch_b})
+        a = batch_a
+        b = batch_b
+
+    assert (
+        sum(
+            isinstance(var.owner.op, (Blockwise | Dot))
+            for var in ancestors([out])
+            if var.owner
+        )
+        == 1
+    )
+
+    # For now rewrite only applies to Batched Dots
+    rewritten_out = rewrite_graph(out)
+    assert rewritten_out.type.shape == out.type.shape
+    assert sum(
+        isinstance(var.owner.op, (Blockwise | Dot))
+        for var in ancestors([rewritten_out])
+        if var.owner
+    ) == (0 if batched else 1)
+
+    a_test = np.random.normal(size=a.type.shape).astype(a.type.dtype)
+    b_test = np.random.normal(size=b.type.shape).astype(b.type.dtype)
+    test_mode = Mode(linker="py", optimizer=None)
+    np.testing.assert_allclose(
+        out.eval({a: a_test, b: b_test}, mode=test_mode),
+        rewritten_out.eval({a: a_test, b: b_test}, mode=test_mode),
+    )

tests/tensor/test_basic.py

@@ -716,6 +716,32 @@ def test_masked_array_not_implemented(
             ptb.as_tensor(x)
 
 
+def check_alloc_runtime_broadcast(mode):
+    """Check we emmit a clear error when runtime broadcasting would occur according to Numpy rules."""
+    floatX = config.floatX
+    x_v = vector("x", shape=(None,))
+
+    out = alloc(x_v, 5, 3)
+    f = pytensor.function([x_v], out, mode=mode)
+    TestAlloc.check_allocs_in_fgraph(f.maker.fgraph, 1)
+
+    np.testing.assert_array_equal(
+        f(x=np.zeros((3,), dtype=floatX)),
+        np.zeros((5, 3), dtype=floatX),
+    )
+    with pytest.raises(ValueError, match="Runtime broadcasting not allowed"):
+        f(x=np.zeros((1,), dtype=floatX))
+
+    out = alloc(specify_shape(x_v, (1,)), 5, 3)
+    f = pytensor.function([x_v], out, mode=mode)
+    TestAlloc.check_allocs_in_fgraph(f.maker.fgraph, 1)
+
+    np.testing.assert_array_equal(
+        f(x=np.zeros((1,), dtype=floatX)),
+        np.zeros((5, 3), dtype=floatX),
+    )
+
+
 class TestAlloc:
     dtype = config.floatX
     mode = mode_opt
@@ -729,32 +755,6 @@ def check_allocs_in_fgraph(fgraph, n):
             == n
         )
 
-    @staticmethod
-    def check_runtime_broadcast(mode):
-        """Check we emmit a clear error when runtime broadcasting would occur according to Numpy rules."""
-        floatX = config.floatX
-        x_v = vector("x", shape=(None,))
-
-        out = alloc(x_v, 5, 3)
-        f = pytensor.function([x_v], out, mode=mode)
-        TestAlloc.check_allocs_in_fgraph(f.maker.fgraph, 1)
-
-        np.testing.assert_array_equal(
-            f(x=np.zeros((3,), dtype=floatX)),
-            np.zeros((5, 3), dtype=floatX),
-        )
-        with pytest.raises(ValueError, match="Runtime broadcasting not allowed"):
-            f(x=np.zeros((1,), dtype=floatX))
-
-        out = alloc(specify_shape(x_v, (1,)), 5, 3)
-        f = pytensor.function([x_v], out, mode=mode)
-        TestAlloc.check_allocs_in_fgraph(f.maker.fgraph, 1)
-
-        np.testing.assert_array_equal(
-            f(x=np.zeros((1,), dtype=floatX)),
-            np.zeros((5, 3), dtype=floatX),
-        )
-
     def setup_method(self):
         self.rng = np.random.default_rng(seed=utt.fetch_seed())
 
@@ -912,7 +912,7 @@ def test_alloc_of_view_linker(self):
 
     @pytest.mark.parametrize("mode", (Mode("py"), Mode("c")))
     def test_runtime_broadcast(self, mode):
-        self.check_runtime_broadcast(mode)
+        check_alloc_runtime_broadcast(mode)
 
 
 def test_infer_static_shape():