Implement BandedDot Op

pytensor/tensor/slinalg.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 import scipy.linalg as scipy_linalg
+from numpy import diag, zeros
 from numpy.exceptions import ComplexWarning
 
 import pytensor
@@ -1669,6 +1670,70 @@ def block_diag(*matrices: TensorVariable):
     return _block_diagonal_matrix(*matrices)
 
 
+class BandedDot(Op):
+    __props__ = ("lower_diags", "upper_diags")
+    gufunc_signature = "(m,n),(n)->(m)"
+
+    def __init__(self, lower_diags, upper_diags):
+        self.lower_diags = lower_diags
+        self.upper_diags = upper_diags
+
+    def make_node(self, A, b):
+        A = as_tensor_variable(A)
+        B = as_tensor_variable(b)
+
+        out_dtype = pytensor.scalar.upcast(A.dtype, B.dtype)
+        output = b.type.clone(dtype=out_dtype)()
+
+        return pytensor.graph.basic.Apply(self, [A, B], [output])
+
+    def infer_shape(self, fgraph, nodes, shapes):
+        return [shapes[0][:-1]]
+
+    def perform(self, node, inputs, outputs_storage):
+        A, b = inputs
+        m, n = A.shape
+
+        kl = self.lower_diags
+        ku = self.upper_diags
+
+        A_banded = zeros((kl + ku + 1, n), dtype=A.dtype, order="C")
+
+        for i, k in enumerate(range(ku, -kl - 1, -1)):
+            A_banded[i, slice(k, None) if k >= 0 else slice(None, n + k)] = diag(A, k=k)
+
+        fn = scipy_linalg.get_blas_funcs("gbmv", dtype=A.dtype)
+        outputs_storage[0][0] = fn(m=m, n=n, kl=kl, ku=ku, alpha=1, a=A_banded, x=b)
+
+
+def banded_dot(A: TensorLike, b: TensorLike, lower_diags: int, upper_diags: int):
+    """
+    Specialized matrix-vector multiplication for cases when A is a banded matrix
+
+    No type-checking is done on A at runtime, so all data in A off the banded diagonals will be ignored. This will lead
+    to incorrect results if A is not actually a banded matrix.
+
+    Unlike dot, this function is only valid if b is a vector.
+
+    Parameters
+    ----------
+    A: Tensorlike
+        Matrix to perform banded dot on.
+    b: Tensorlike
+        Vector to perform banded dot on.
+    lower_diags: int
+        Number of nonzero lower diagonals of A
+    upper_diags: int
+        Number of nonzero upper diagonals of A
+
+    Returns
+    -------
+    out: Tensor
+        The matrix multiplication result
+    """
+    return Blockwise(BandedDot(lower_diags, upper_diags))(A, b)
+
+
 __all__ = [
     "cholesky",
     "solve",
@@ -1683,4 +1748,5 @@ def block_diag(*matrices: TensorVariable):
     "lu",
     "lu_factor",
     "lu_solve",
+    "banded_dot",
 ]

tests/tensor/test_slinalg.py

@@ -12,11 +12,13 @@
 from pytensor.graph.basic import equal_computations
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.slinalg import (
+    BandedDot,
     Cholesky,
     CholeskySolve,
     Solve,
     SolveBase,
     SolveTriangular,
+    banded_dot,
     block_diag,
     cho_solve,
     cholesky,
@@ -1051,3 +1053,63 @@ def test_block_diagonal_blockwise():
     B = np.random.normal(size=(1, batch_size, 4, 4)).astype(config.floatX)
     result = block_diag(A, B).eval()
     assert result.shape == (10, batch_size, 6, 6)
+
+
+def _make_banded_A(A, kl, ku):
+    diag_idxs = range(-kl, ku + 1)
+    diags = (np.diag(A, k=k) for k in diag_idxs)
+    return sum(np.diag(d, k=k) for k, d in zip(diag_idxs, diags))
+
+
+@pytest.mark.parametrize(
+    "A_shape",
+    [
+        (10, 10),
+    ],
+)
+@pytest.mark.parametrize(
+    "kl, ku", [(1, 1), (0, 1), (2, 2)], ids=["tridiag", "upper-only", "banded"]
+)
+def test_banded_dot(A_shape, kl, ku):
+    rng = np.random.default_rng()
+
+    A_val = _make_banded_A(rng.normal(size=A_shape), kl=kl, ku=ku).astype(config.floatX)
+    b_val = rng.normal(size=(A_shape[-1],)).astype(config.floatX)
+
+    A = pt.tensor("A", shape=A_val.shape, dtype=A_val.dtype)
+    b = pt.tensor("b", shape=b_val.shape, dtype=b_val.dtype)
+    res = banded_dot(A, b, kl, ku)
+    res_2 = A @ b
+
+    fn = function([A, b], [res, res_2], trust_input=True)
+    assert any(isinstance(node.op, BandedDot) for node in fn.maker.fgraph.apply_nodes)
+
+    x_val, x2_val = fn(A_val, b_val)
+
+    np.testing.assert_allclose(x_val, x2_val)
+
+
+@pytest.mark.parametrize("op", ["dot", "banded_dot"], ids=str)
+@pytest.mark.parametrize(
+    "A_shape",
+    [(10, 10), (100, 100), (1000, 1000), (10_000, 10_000)],
+    ids=["10", "100", "1000", "10_000"],
+)
+def test_banded_dot_perf(op, A_shape, benchmark):
+    rng = np.random.default_rng()
+
+    A_val = _make_banded_A(rng.normal(size=A_shape), kl=1, ku=1).astype(config.floatX)
+    b_val = rng.normal(size=(A_shape[-1],)).astype(config.floatX)
+
+    A = pt.tensor("A", shape=A_val.shape, dtype=A_val.dtype)
+    b = pt.tensor("b", shape=b_val.shape, dtype=A_val.dtype)
+
+    if op == "dot":
+        f = pt.dot
+    elif op == "banded_dot":
+        f = functools.partial(banded_dot, lower_diags=1, upper_diags=1)
+
+    res = f(A, b)
+    fn = function([A, b], res, trust_input=True)
+
+    benchmark(fn, A_val, b_val)