Add rewrite for diag(A @ B) to avoid materializing the product

pytensor/tensor/rewriting/subtensor.py

@@ -192,6 +192,54 @@ def _arange_provably_unique(start, stop, step) -> bool:
     return False
 
 
+def _arange_shifted_bounds(idx):
+    """``arange(a, b, s) +/- k`` as the bounds of the ``arange`` it is equal to.
+
+    A scalar shift slides the whole range, so whether the entries straddle zero — the
+    only thing that makes distinct values alias onto one position — is still decided by
+    the bounds, once the shift is folded into them. ``arange(-5, 5) + 5`` is
+    ``arange(0, 10)``, unique, even though the unshifted range is not.
+    """
+    if not (
+        isinstance(idx.owner_op, Elemwise)
+        and isinstance(idx.owner.op.scalar_op, Add | Sub)
+    ):
+        return None
+    x, y = idx.owner.inputs
+    if isinstance(x.owner_op, ARange) and all(y.type.broadcastable):
+        arange_var, shift_var = x, y
+    elif (
+        isinstance(idx.owner.op.scalar_op, Add)
+        and isinstance(y.owner_op, ARange)
+        and all(x.type.broadcastable)
+    ):
+        arange_var, shift_var = y, x
+    else:
+        return None
+    try:
+        shift = int(get_underlying_scalar_constant_value(shift_var))
+    except NotScalarConstantError:
+        # A symbolic shift moves the range by an unknown amount, so where it sits
+        # relative to zero -- the only thing that matters here -- stays unknown.
+        return None
+    if isinstance(idx.owner.op.scalar_op, Sub):
+        shift = -shift
+
+    def shifted(bound):
+        # Fold eagerly: the sign checks below read the bound itself, not a value
+        # they could recover from an unevaluated Add. The shifted value can leave the
+        # original bound's dtype (``arange(5)`` carries int8 bounds), and only its
+        # sign and magnitude are ever read, so let the constant pick its own.
+        try:
+            value = int(get_scalar_constant_value(bound)) + shift
+        except NotScalarConstantError:
+            return bound + shift
+        return tensor_constant(np.asarray(value))
+
+    start, stop, step = arange_var.owner.inputs
+    return shifted(start), shifted(stop), step
+
+
 def _index_provably_unique(idx, fgraph: FunctionGraph | None) -> bool:
     """Whether a single index selects each position on its own axis at most once.
 
@@ -217,6 +265,8 @@ def _index_provably_unique(idx, fgraph: FunctionGraph | None) -> bool:
         return True
     if isinstance(idx.owner_op, ARange):
         return _arange_provably_unique(*idx.owner.inputs)
+    if (shifted_bounds := _arange_shifted_bounds(idx)) is not None:
+        return _arange_provably_unique(*shifted_bounds)
     if isinstance(idx.owner_op, Reshape | DimShuffle):
         # Views that only reorder or insert size-1 dims keep the value multiset.
         return _index_provably_unique(idx.owner.inputs[0], fgraph)
@@ -384,7 +434,8 @@ def eager_add_zero(x, y):
 def _eager_scalar(x):
     """Reduce a 0d or ``(1,)``-shaped tensor to the simplest scalar form."""
     if isinstance(x, TensorConstant):
-        return int(x.data)
+        # ``.item()`` covers the ``(1,)`` case the 0d-only ``int()`` rejects.
+        return int(x.data.item())
     if isinstance(x.owner_op, DimShuffle) and x.owner.op.input_ndim == 0:
         inner = x.owner.inputs[0]
         return int(inner.data) if isinstance(inner, TensorConstant) else inner

pytensor/tensor/rewriting/subtensor_lift.py

@@ -344,6 +344,94 @@ def local_subtensor_of_dot(fgraph, node):
     return [r]
 
 
+@register_canonicalize
+@register_stabilize
+@register_specialize
+@node_rewriter([AdvancedSubtensor])
+def local_advanced_subtensor_of_dot(fgraph, node):
+    """Push an advanced index of a matrix product into its operands.
+
+    ``dot(A, B)[i]    -> dot(A[i], B)``
+    ``dot(A, B)[:, j] -> dot(A, B[:, j])``
+    ``dot(A, B)[i, j] -> (A[i] * B.mT[j]).sum(-1)``
+
+    Since ``dot(A, B)[i, j] = sum_c A[i, c] * B[c, j]``, an index on one product axis
+    only gathers from the operand that axis comes from, leaving a smaller ``Dot``.
+    Indexing *both* axes couples them, and a coupled index cannot be a product at all
+    — its output index appears on both operands (``ik,ki->i``) — so it becomes a
+    gather of rows of ``A`` and of ``B.mT`` contracted elementwise.
+    ``ExtractDiag`` is the ``i is j is arange(d)`` case, which
+    ``local_extract_diag_lift`` lowers into this form.
+
+    This is the advanced-indexing counterpart of ``local_subtensor_of_dot``, whose
+    basic indices can only ever be separable.
+    """
+    x, *idx_vars = node.inputs
+
+    match x.owner_op:
+        case Dot() | Blockwise(Dot()):
+            pass
+        case _:
+            return None
+
+    # If the product feeds other clients it is materialized anyway, so pushing the
+    # index into the operands would only add work.
+    if len(fgraph.clients[x]) > 1:
+        return None
+
+    ndim = x.type.ndim
+    idx_tuple = indices_from_subtensor(idx_vars, node.op.idx_list)
+    idx_tuple = (*idx_tuple, *(slice(None),) * (ndim - len(idx_tuple)))
+    *batch_idx, row_idx, col_idx = idx_tuple
+    if any(idx != slice(None) for idx in batch_idx):
+        return None
+
+    def is_advanced(idx):
+        # Boolean masks have a data-dependent shape the size gate can't weigh.
+        return (
+            isinstance(idx, TensorVariable)
+            and idx.type.ndim > 0
+            and idx.type.dtype != "bool"
+        )
+
+    row_adv, col_adv = is_advanced(row_idx), is_advanced(col_idx)
+    row_dim, col_dim = x.type.shape[-2:]
+    a, b = x.owner.inputs
+    a_batch = (slice(None),) * (a.type.ndim - 2)
+    b_batch = (slice(None),) * (b.type.ndim - 2)
+
+    if row_adv and col_adv:
+        if not _indices_provably_not_larger(
+            [(row_idx, row_dim), (col_idx, col_dim)], fgraph
+        ):
+            return None
+        a_rows = a[(*a_batch, row_idx)]
+        b_rows = b.mT[(*b_batch, col_idx)]
+        new_out = (a_rows * b_rows).sum(-1)
+    elif row_adv and col_idx == slice(None):
+        # A multidimensional index would leave the gathered operand with more than
+        # the two core dims `matmul` contracts, so only vector indices are pushed.
+        if row_idx.type.ndim != 1:
+            return None
+        if not _indices_provably_not_larger([(row_idx, row_dim)], fgraph):
+            return None
+        new_out = a[(*a_batch, row_idx)] @ b
+    elif col_adv and row_idx == slice(None):
+        if col_idx.type.ndim != 1:
+            return None
+        if not _indices_provably_not_larger([(col_idx, col_dim)], fgraph):
+            return None
+        new_out = a @ b[(*b_batch, slice(None), col_idx)]
+    else:
+        return None
+
+    if new_out.type.dtype != node.outputs[0].type.dtype:
+        new_out = new_out.astype(node.outputs[0].type.dtype)
+
+    copy_stack_trace(node.outputs[0], new_out)
+    return [new_out]
+
+
 @register_canonicalize("shape_unsafe")
 @register_specialize("shape_unsafe")
 @node_rewriter([Subtensor, AdvancedSubtensor])
@@ -903,9 +991,13 @@ def local_extract_diag_lift(fgraph, node):
     - ``Alloc`` parent — gated on ``local_subtensor_of_alloc``; rebuilds the
       outer Alloc shape with the diag length ``d`` so ``Shape(arange(d))[0]``
       doesn't survive.
+    - ``Dot`` / ``Blockwise(Dot)`` parent over the two contracted axes — gated on
+      ``local_advanced_subtensor_of_dot``, which turns the paired-arange gather
+      into ``(A * B.mT).sum(-1)``.
     - ``Elemwise`` / ``Blockwise`` parent — gated on
       ``local_subtensor_of_batch_dims``. Blockwise core dims bail (the lift
-      can't push past core ndim).
+      can't push past core ndim), so a batched dot diagonalized over its batch
+      dims lands here rather than on the branch above.
     - ``AdvancedIncSubtensor`` write-chain parent — gated on
       ``local_advanced_read_of_write_constant_indices``.
     """
@@ -916,14 +1008,25 @@ def local_extract_diag_lift(fgraph, node):
 
     if isinstance(parent_op, Alloc):
         shape_inputs = inner.owner.inputs[1:]
+        gate = local_subtensor_of_alloc
+    # A dot whose diagonal straddles both core dims: the paired index couples the
+    # two operands, which only the fold below can express.
+    elif (
+        isinstance(parent_op, Dot)
+        or (isinstance(parent_op, Blockwise) and isinstance(parent_op.core_op, Dot))
+    ) and (a1, a2) == (inner.type.ndim - 2, inner.type.ndim - 1):
+        shape_inputs = inner.shape
+        gate = local_advanced_subtensor_of_dot
     elif isinstance(parent_op, Elemwise | Blockwise):
         if isinstance(parent_op, Blockwise):
             batch_ndim = inner.owner.op.batch_ndim(inner.owner)
             if a1 >= batch_ndim or a2 >= batch_ndim:
                 return None
         shape_inputs = inner.shape
+        gate = local_subtensor_of_batch_dims
     elif isinstance(parent_op, AdvancedIncSubtensor):
         shape_inputs = inner.shape
+        gate = local_advanced_read_of_write_constant_indices
     else:
         return None
 
@@ -943,13 +1046,7 @@ def _diag_dim(ax, off):
     idxs = _diag_indices(inner.type.ndim, a1, a2, diag_len, row_off, col_off)
     hypothetical = inner[tuple(idxs)].owner
 
-    if isinstance(parent_op, Alloc):
-        return local_subtensor_of_alloc.fn(fgraph, hypothetical)
-    elif isinstance(parent_op, Elemwise | Blockwise):
-        return local_subtensor_of_batch_dims.fn(fgraph, hypothetical)
-    else:
-        # AdvancedIncSubtensor write chain
-        return local_advanced_read_of_write_constant_indices.fn(fgraph, hypothetical)
+    return gate.fn(fgraph, hypothetical)
 
 
 extract_diag_lift_pass = SequentialGraphRewriter(
@@ -958,6 +1055,7 @@ def _diag_dim(ax, off):
         local_extract_diag_of_eye,
         local_extract_diag_lift,
         local_subtensor_of_alloc,
+        local_advanced_subtensor_of_dot,
         local_subtensor_of_batch_dims,
         local_slice_read_of_write,
         local_advanced_read_of_write_constant_indices,

tests/tensor/rewriting/linalg/test_products.py

@@ -9,12 +9,13 @@
 from pytensor.configdefaults import config
 from pytensor.graph import FunctionGraph, ancestors
 from pytensor.graph.rewriting.utils import rewrite_graph
-from pytensor.tensor.basic import alloc_diag
+from pytensor.tensor.basic import ExtractDiag, alloc_diag
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.linalg.constructors import BlockDiagonal
 from pytensor.tensor.linalg.products import KroneckerProduct
 from pytensor.tensor.math import Dot
-from tests.unittest_tools import assert_equal_computations
+from pytensor.tensor.rewriting.subtensor_lift import extract_diag_lift_pass
+from tests.unittest_tools import RewriteTester, assert_equal_computations
 
 
 def test_nested_blockdiag_fusion():
@@ -482,3 +483,128 @@ def test_det_of_non_permutation_selection_is_not_rewritten():
     f = function([idx], pt.linalg.det(S), mode="FAST_RUN")
     assert any(isinstance(node.op, Det) for node in f.maker.fgraph.apply_nodes)
     assert_allclose(f(np.array([0, 1, 1, 3])), 0.0)  # duplicate column -> singular
+
+
+class TestDiagOfDot:
+    """``diag(A @ B)`` is lowered to a paired-arange gather by
+    ``local_extract_diag_lift`` and folded into the operands by
+    ``local_advanced_subtensor_of_dot``, so it needs the specialize pass.
+    """
+
+    rewrite_kw = dict(include=("ShapeOpt", "canonicalize", "specialize"))
+
+    @pytest.mark.parametrize("batch", [(), (3,)], ids=["2d", "batched"])
+    @pytest.mark.parametrize("offset", [0, 1, -1, 2, -2], ids=str)
+    def test_diag_of_dot(self, batch, offset):
+        rng = np.random.default_rng(0)
+        n, contracted = 6, 4
+        A = pt.tensor("A", shape=(*batch, n, contracted))
+        B = pt.tensor("B", shape=(*batch, contracted, n))
+
+        out = pt.diagonal(A @ B, offset=offset, axis1=-2, axis2=-1)
+        result = RewriteTester([A, B], [out], **self.rewrite_kw)
+
+        # The main diagonal has a stable closed form; shifted diagonals only differ
+        # by slicing, whose exact canonical shape is verified against the oracle below.
+        if offset == 0:
+            result.assert_graph((A * B.mT).sum(-1))
+
+        result.assert_eval(
+            rng.standard_normal((*batch, n, contracted)),
+            rng.standard_normal((*batch, contracted, n)),
+        )
+
+    @pytest.mark.parametrize("offset", [9, -9], ids=str)
+    def test_out_of_range_offset(self, offset):
+        """An offset past the matrix gives an empty diagonal, so only values are checked."""
+        rng = np.random.default_rng(1)
+        A = pt.matrix("A", shape=(6, 6))
+        B = pt.matrix("B", shape=(6, 6))
+        result = RewriteTester(
+            [A, B], [pt.diagonal(A @ B, offset=offset)], **self.rewrite_kw
+        )
+        result.assert_eval(rng.standard_normal((6, 6)), rng.standard_normal((6, 6)))
+
+    def test_rectangular_product(self):
+        """A non-square product exercises the trim to the shorter side."""
+        rng = np.random.default_rng(2)
+        A = pt.matrix("A", shape=(6, 4))
+        B = pt.matrix("B", shape=(4, 3))
+        result = RewriteTester([A, B], [pt.diag(A @ B)], **self.rewrite_kw)
+        result.assert_eval(rng.standard_normal((6, 4)), rng.standard_normal((4, 3)))
+
+    @pytest.mark.parametrize("offset", [0, 2, -2], ids=str)
+    def test_dynamic_shapes(self, offset):
+        """With no static shapes the diagonal length stays symbolic, so the size gate
+        rests on the paired ``arange``. An offset diagonal indexes with ``arange(d) +
+        k``, whose entry count is the ``arange``'s even though the shift can alias
+        entries onto each other -- the gate has to read it as a size, not a uniqueness.
+        """
+        rng = np.random.default_rng(3)
+        A = pt.matrix("A")
+        B = pt.matrix("B")
+        result = RewriteTester(
+            [A, B], [pt.diagonal(A @ B, offset=offset)], **self.rewrite_kw
+        )
+        assert not any(
+            isinstance(node.op, Dot | ExtractDiag)
+            for node in result.rewr_fg.apply_nodes
+        )
+        result.assert_eval(rng.standard_normal((6, 4)), rng.standard_normal((4, 6)))
+
+    def test_trace_of_dot(self):
+        """``trace`` is ``diagonal(...).sum()``, so it folds through the same path."""
+        rng = np.random.default_rng(4)
+        A = pt.matrix("A", shape=(6, 4))
+        B = pt.matrix("B", shape=(4, 6))
+        result = RewriteTester([A, B], [pt.trace(A @ B)], **self.rewrite_kw)
+        assert not any(isinstance(node.op, Dot) for node in result.rewr_fg.apply_nodes)
+        result.assert_eval(rng.standard_normal((6, 4)), rng.standard_normal((4, 6)))
+
+    def test_diag_of_einsum(self):
+        """``einsum`` only inlines to a ``Dot`` during specialize, so the fold has
+        to be reachable there and not only in canonicalize."""
+        rng = np.random.default_rng(5)
+        A = pt.matrix("A", shape=(6, 4))
+        B = pt.matrix("B", shape=(4, 6))
+        out = pt.diagonal(pt.einsum("ik,kj->ij", A, B))
+        result = RewriteTester([A, B], [out], **self.rewrite_kw)
+        assert not any(isinstance(node.op, Dot) for node in result.rewr_fg.apply_nodes)
+        result.assert_eval(rng.standard_normal((6, 4)), rng.standard_normal((4, 6)))
+
+    def test_single_client_guard(self):
+        """The product also feeds ``M.sum()`` so it is materialized regardless; the
+        rewrite must not fire and add a redundant elementwise path."""
+        A = pt.matrix("A", shape=(5, 5))
+        B = pt.matrix("B", shape=(5, 5))
+        M = pt.dot(A, B)
+        result = RewriteTester(
+            [A, B],
+            [pt.diag(M) + M.sum()],
+            include=None,
+            custom_rewrite=extract_diag_lift_pass,
+        )
+        result.assert_graph(pt.diag(M) + M.sum())
+
+    def test_wrong_axis_guard(self):
+        """A diagonal over the batch/row axes, not the contracted axes, has no
+        closed form in the operands and must be left alone."""
+        A = pt.tensor("A", shape=(4, 5, 5))
+        B = pt.tensor("B", shape=(4, 5, 5))
+        out = pt.diagonal(A @ B, axis1=0, axis2=1)
+        result = RewriteTester([A, B], [out], **self.rewrite_kw)
+        assert any(
+            isinstance(node.op, ExtractDiag) for node in result.rewr_fg.apply_nodes
+        )
+
+    def test_dtype_preserved(self):
+        """``sum`` upcasts an int32 accumulator to int64, so the rewrite must cast back."""
+        rng = np.random.default_rng(6)
+        A = pt.matrix("A", shape=(6, 4), dtype="int32")
+        B = pt.matrix("B", shape=(4, 6), dtype="int32")
+        result = RewriteTester([A, B], [pt.diag(A @ B)], **self.rewrite_kw)
+        result.assert_graph((A * B.mT).sum(-1).astype("int32"))
+        result.assert_eval(
+            rng.integers(0, 5, (6, 4)).astype("int32"),
+            rng.integers(0, 5, (4, 6)).astype("int32"),
+        )