Add numba overload for pytensor.tensor.slinalg.block_diag

Author: jessegrabowski

pytensor/link/numba/dispatch/slinalg.py

@@ -9,7 +9,7 @@
 
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import numba_funcify
-from pytensor.tensor.slinalg import SolveTriangular
+from pytensor.tensor.slinalg import BlockDiagonalMatrix, SolveTriangular
 
 
 _PTR = ctypes.POINTER
@@ -273,3 +273,25 @@ def solve_triangular(a, b):
         return res
 
     return solve_triangular
+
+
+@numba_funcify.register(BlockDiagonalMatrix)
+def numba_funcify_BlockDiagonalMatrix(op, node, **kwargs):
+    dtype = node.outputs[0].dtype
+
+    # TODO: Why do we always inline all functions? It doesn't work with starred args, so can't use it in this case.
+    @numba_basic.numba_njit(inline="never")
+    def block_diag(*arrs):
+        shapes = np.array([a.shape for a in arrs], dtype=dtype)
+        out_shape = [int(s) for s in np.sum(shapes, axis=0)]
+        out = np.zeros((out_shape[0], out_shape[1]))
+
+        r, c = 0, 0
+        for arr, shape in zip(arrs, shapes):
+            rr, cc = shape
+            out[r : r + rr, c : c + cc] = arr
+            r += rr
+            c += cc
+        return out
+
+    return block_diag

pytensor/tensor/slinalg.py

@@ -949,7 +949,7 @@ def make_node(self, *matrices, name=None):
         return Apply(self, matrices, [out_type])
 
     def perform(self, node, inputs, output_storage, params=None):
-        dtype = largest_common_dtype(inputs)
+        dtype = node.outputs[0].type.dtype
         output_storage[0][0] = scipy.linalg.block_diag(*inputs).astype(dtype)
 
 

tests/link/numba/test_slinalg.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+from scipy import linalg
 
 import pytensor
 import pytensor.tensor as pt
@@ -10,7 +11,6 @@
 
 numba = pytest.importorskip("numba")
 
-
 ATOL = 0 if config.floatX.endswith("64") else 1e-6
 RTOL = 1e-7 if config.floatX.endswith("64") else 1e-6
 rng = np.random.default_rng(42849)
@@ -102,3 +102,30 @@ def test_solve_triangular_raises_on_nan_inf(value):
         ValueError, match=re.escape("Non-numeric values (nan or inf) returned ")
     ):
         f(A_tri, b)
+
+
+def test_block_diag():
+    A = pt.matrix("A")
+    B = pt.matrix("B")
+    C = pt.matrix("C")
+    D = pt.matrix("D")
+    X = pt.linalg.block_diag(A, B, C, D)
+    f = pytensor.function([A, B, C, D], X, mode="NUMBA")
+
+    A_val = np.random.normal(size=(5, 5))
+    B_val = np.random.normal(size=(3, 3))
+    C_val = np.random.normal(size=(2, 2))
+    D_val = np.random.normal(size=(4, 4))
+
+    X_val = f(A_val, B_val, C_val, D_val)
+    np.testing.assert_allclose(
+        np.block([[A_val, np.zeros((5, 3))], [np.zeros((3, 5)), B_val]]), X_val[:8, :8]
+    )
+    np.testing.assert_allclose(
+        np.block([[C_val, np.zeros((2, 4))], [np.zeros((4, 2)), D_val]]), X_val[8:, 8:]
+    )
+    np.testing.assert_allclose(np.zeros((8, 6)), X_val[:8, 8:])
+    np.testing.assert_allclose(np.zeros((6, 8)), X_val[8:, :8])
+
+    X_sp = linalg.block_diag(A_val, B_val, C_val, D_val)
+    np.testing.assert_allclose(X_val, X_sp)

tests/tensor/test_slinalg.py

@@ -665,6 +665,7 @@ def test_solve_discrete_are_grad():
 
 
 def test_block_diagonal():
-    matrices = [np.array([[1.0, 2.0], [3.0, 4.0]]), np.array([[5.0, 6.0], [7.0, 8.0]])]
-    result = block_diag(*matrices)
-    np.testing.assert_allclose(result.eval(), scipy.linalg.block_diag(*matrices))
+    A = np.array([[1.0, 2.0], [3.0, 4.0]])
+    B = np.array([[5.0, 6.0], [7.0, 8.0]])
+    result = block_diag(A, B, name="X")
+    np.testing.assert_allclose(result.eval(), scipy.linalg.block_diag(A, B))