add jax overload for pytensor.tensor.slinalg.block_diag

Author: jessegrabowski

pytensor/link/jax/dispatch/slinalg.py

@@ -1,7 +1,12 @@
 import jax
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
-from pytensor.tensor.slinalg import Cholesky, Solve, SolveTriangular
+from pytensor.tensor.slinalg import (
+    BlockDiagonalMatrix,
+    Cholesky,
+    Solve,
+    SolveTriangular,
+)
 
 
 @jax_funcify.register(Cholesky)
@@ -45,3 +50,11 @@ def solve_triangular(A, b):
         )
 
     return solve_triangular
+
+
+@jax_funcify.register(BlockDiagonalMatrix)
+def jax_funcify_BlockDiagonalMatrix(op, **kwargs):
+    def block_diag(*inputs):
+        return jax.scipy.linalg.block_diag(*inputs)
+
+    return block_diag

tests/link/jax/test_slinalg.py

@@ -129,3 +129,23 @@ def test_jax_SolveTriangular(trans, lower, check_finite):
             np.arange(10).astype(config.floatX),
         ],
     )
+
+
+def test_jax_block_diag():
+    A = matrix("A")
+    B = matrix("B")
+    C = matrix("C")
+    D = matrix("D")
+
+    out = pt_slinalg.block_diag(A, B, C, D)
+
+    out_fg = FunctionGraph([A, B, C, D], [out])
+    compare_jax_and_py(
+        out_fg,
+        [
+            np.random.normal(size=(5, 5)).astype(config.floatX),
+            np.random.normal(size=(3, 3)).astype(config.floatX),
+            np.random.normal(size=(2, 2)).astype(config.floatX),
+            np.random.normal(size=(4, 4)).astype(config.floatX),
+        ],
+    )