Repository: google/jax

Issue:
Better support for handling block diagonal matrices
Dear JAX team,

Thanks for the amazing work with JAX!

I'm currently doing some work which would really benefit from support for handling block diagonal matrices. In particular, I would like to assemble matrices using something like `scipy.linalg.block_diag` (https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.linalg.block_diag.html). It would also be nice to be able to split them again, but that's a bit awkward in numpy too and I think could be done using slicing. I've been able to work around this for now using kronecker products, but it's not very neat.

Would this be hard to add? Could give it a go too if it's not too hard and there's a rough plan (maybe scatter would be one way to go?).


Discussion:
https://github.com/google/jax/pull/2106 adds an implementation of `np.block`, which might help with assembling block matrices.
How large are your blocks?

To me, this looks like a usecase that might be better solved by a `LinearOperator` object that supports (at least) matrix-vector multiplication, e.g., https://www.tensorflow.org/api_docs/python/tf/linalg/LinearOperatorBlockDiag
Thanks both of you! 

@hawkinsp I think np.block will do the trick! It doesn't look quite as nice as block_diag because you have to explicitly create the zero parts, but that's not a big deal. Happy to have the issue closed!

@shoyer The blocks are currently really small, definitely under 10x10. For larger ones I agree, the linear operator approach should be much more efficient. Are there any plans to implement them in JAX?
> Are there any plans to implement them in JAX?

Not that I know of, currently

Resolving patch:
diff --git a/jax/scipy/linalg.py b/jax/scipy/linalg.py
--- a/jax/scipy/linalg.py
+++ b/jax/scipy/linalg.py
@@ -278,7 +278,7 @@ def _expm(A, upper_triangular=False):
     R = _solve_P_Q(P, Q, upper_triangular)
     R = _squaring(R, n_squarings)
     return R
-    
+
 @jit
 def _calc_P_Q(A):
     A = np.asarray(A)
@@ -379,3 +379,25 @@ def _pade13(A):
     U = np.dot(A,np.dot(A6, b[13]*A6 + b[11]*A4 + b[9]*A2) + b[7]*A6 + b[5]*A4 + b[3]*A2 + b[1]*ident)
     V = np.dot(A6, b[12]*A6 + b[10]*A4 + b[8]*A2) + b[6]*A6 + b[4]*A4 + b[2]*A2 + b[0]*ident
     return U,V
+
+
+@_wraps(scipy.linalg.block_diag)
+@jit
+def block_diag(*arrs):
+  if len(arrs) == 0:
+    arrs = [np.zeros((1, 0))]
+  arrs = np._promote_dtypes(*arrs)
+  bad_shapes = [i for i, a in enumerate(arrs) if np.ndim(a) > 2]
+  if bad_shapes:
+    raise ValueError("Arguments to jax.scipy.linalg.block_diag must have at "
+                     "most 2 dimensions, got {} at argument {}."
+                     .format(arrs[bad_shapes[0]], bad_shapes[0]))
+  arrs = [np.atleast_2d(a) for a in arrs]
+  acc = arrs[0]
+  dtype = lax.dtype(acc)
+  for a in arrs[1:]:
+    _, c = a.shape
+    a = lax.pad(a, dtype.type(0), ((0, 0, 0), (acc.shape[-1], 0, 0)))
+    acc = lax.pad(acc, dtype.type(0), ((0, 0, 0), (0, c, 0)))
+    acc = lax.concatenate([acc, a], dimension=0)
+  return acc


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -682,6 +682,23 @@ def f(inp):
 
 class ScipyLinalgTest(jtu.JaxTestCase):
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_i={}".format(i), "args": args}
+      for i, args in enumerate([
+        (),
+        (1,),
+        (7, -2),
+        (3, 4, 5),
+        (onp.ones((3, 4), dtype=np.float_), 5,
+         onp.random.randn(5, 2).astype(np.float_)),
+      ])))
+  def testBlockDiag(self, args):
+    args_maker = lambda: args
+    self._CheckAgainstNumpy(osp.linalg.block_diag, jsp.linalg.block_diag,
+                            args_maker, check_dtypes=True)
+    self._CompileAndCheck(jsp.linalg.block_diag, args_maker, check_dtypes=True)
+
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name":
        "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),