Repository: google/jax
Issue: linalg.solve doesn't respect broadcasting semantics

I'm attempting a batched matrix solve,
```python
import jax.numpy as jp
import numpy as np

A = np.eye(2)[np.newaxis]
b = np.ones((3, 2))
print(A.shape, b.shape)
print(np.linalg.solve(A, b))
print(jp.linalg.solve(A, b))
```
but I'm getting the following:
```
(1, 2, 2) (3, 2)
[[1. 1.]
 [1. 1.]
 [1. 1.]]
/Users/skainswo/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/mle_normal.py in <module>
     31 print(A.shape, b.shape)
     32 print(np.linalg.solve(A, b))
---> 33 print(jp.linalg.solve(A, b))

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/numpy/linalg.py in solve(a, b)
    247   x = b if a_ndims == b_ndims else b[..., None]
    248 
--> 249   permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
    250   x = x[..., permutation, :]
    251 

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax_linalg.py in lu_pivots_to_permutation(swaps, k)
    462     return swaps, permutation
    463 
--> 464   n, = np.shape(swaps)
    465   permutation = np.arange(k)
    466   _, permutation = lax.fori_loop(

ValueError: too many values to unpack (expected 1)
```

I'm on jax 0.1.36 and numpy 1.16.4.

Changed files:
tests/linalg_test.py, jax/lax/lax.py, jax/numpy/linalg.py, jax/lax_linalg.py, tests/lax_numpy_indexing_test.py

diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 3dd0bd4e5..0f065566d 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1164,7 +1164,7 @@ def index_in_dim(operand, index, axis=0, keepdims=True):
 
 def dynamic_slice_in_dim(operand, start_index, slice_size, axis=0):
   """Convenience wrapper around dynamic_slice applying to one dimension."""
-  start_indices = [onp.array([0])] * operand.ndim
+  start_indices = [onp.array([0], dtype=_dtype(start_index))] * operand.ndim
   slice_sizes = list(operand.shape)
 
   axis = int(axis)
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
index e44cdd243..4a3c7d9ec 100644
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -19,9 +19,10 @@ from __future__ import print_function
 import numpy as onp
 
 from jax.numpy import lax_numpy as np
+from jax import ad_util
 from jax import core
 from jax import lax
-from jax import ad_util
+from jax import ops
 from jax.interpreters import xla
 from jax.interpreters import ad
 from jax.interpreters import batching
@@ -450,23 +451,34 @@ def lu_cpu_translation_rule(c, operand):
 xla.backend_specific_translations['cpu'][lu_p] = lu_cpu_translation_rule
 
 
-def lu_pivots_to_permutation(swaps, k):
-  """Converts the pivots (row swaps) returned by LU to a permutation."""
-
-  def body_fn(i, loop_carry):
-    swaps, permutation = loop_carry
-    j = swaps[i]
-    x, y = np.ravel(permutation[i]), np.ravel(permutation[j])
-    permutation = lax.dynamic_update_index_in_dim(permutation, y, i, axis=0)
-    permutation = lax.dynamic_update_index_in_dim(permutation, x, j, axis=0)
-    return swaps, permutation
-
-  n, = np.shape(swaps)
-  permutation = np.arange(k)
-  _, permutation = lax.fori_loop(
-      onp.array(0, onp.int32), onp.array(n, onp.int32), body_fn, (swaps, permutation))
-  return permutation
-
+def lu_pivots_to_permutation(swaps, m):
+  """Converts the pivots (row swaps) returned by LU to a permutation.
+
+  We build a permutation rather than applying `swaps` directly to the rows
+  of a matrix because lax loops aren't differentiable.
+
+  Args:
+    swaps: an array of shape (..., k) of row swaps to perform
+    m: the size of the output permutation. m should be >= k.
+  Returns:
+    An int32 array of shape (..., m).
+  """
+  assert len(swaps.shape) >= 1
+  batch_dims = swaps.shape[:-1]
+  k = swaps.shape[-1]
+
+  def body_fn(i, permutation):
+    j = swaps[..., i]
+    iotas = np.ix_(*(lax.iota(np.int32, b) for b in batch_dims))
+    x = permutation[..., i]
+    y = permutation[iotas + (j,)]
+    permutation = ops.index_update(permutation, ops.index[..., i], y)
+    return ops.index_update(permutation, ops.index[iotas + (j,)], x)
+
+  permutation = lax.broadcasted_iota(np.int32, batch_dims + (m,),
+                                     len(batch_dims))
+  return lax.fori_loop(
+    onp.array(0, onp.int32), onp.array(k, onp.int32), body_fn, permutation)
 
 
 # QR decomposition
diff --git a/jax/numpy/linalg.py b/jax/numpy/linalg.py
index 038fe4962..608d5f522 100644
--- a/jax/numpy/linalg.py
+++ b/jax/numpy/linalg.py
@@ -229,30 +229,35 @@ def solve(a, b):
   b_shape = np.shape(b)
   a_ndims = len(a_shape)
   b_ndims = len(b_shape)
-  if not (a_ndims >= 2 and a_shape[-1] == a_shape[-2] and
-          (a_ndims == b_ndims or a_ndims == b_ndims + 1)):
+  if not (a_ndims >= 2 and a_shape[-1] == a_shape[-2] and b_ndims >= 1):
     msg = ("The arguments to solve must have shapes a=[..., m, m] and "
            "b=[..., m, k] or b=[..., m]; got a={} and b={}")
     raise ValueError(msg.format(a_shape, b_shape))
   lu, pivots = lax_linalg.lu(a)
   dtype = lax.dtype(a)
 
-  # TODO(phawkins): add unit_diagonal support to solve_triangular, use it here
-  # instead of explicit masking of l.
   m = a_shape[-1]
-  l = np.tril(lu, -1)[:, :m] + np.eye(m, m, dtype=dtype)
 
-  # TODO(phawkins): triangular_solve only supports matrices on the RHS, so we
-  # add a dummy dimension. Extend it to support vectors and simplify this.
-  x = b if a_ndims == b_ndims else b[..., None]
+  # Numpy treats the RHS as a (batched) vector if the number of dimensions
+  # differ by 1. Otherwise, broadcasting rules apply.
+  x = b[..., None] if a_ndims == b_ndims + 1 else b
+
+  batch_dims = lax.broadcast_shapes(lu.shape[:-2], x.shape[:-2])
+  x = np.broadcast_to(x, batch_dims + x.shape[-2:])
+  lu = np.broadcast_to(lu, batch_dims + lu.shape[-2:])
 
   permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
-  x = x[..., permutation, :]
+  permutation = np.broadcast_to(permutation, batch_dims + (m,))
+  iotas = np.ix_(*(lax.iota(np.int32, b) for b in batch_dims + (1,)))
+  x = x[iotas[:-1] + (permutation, slice(None))]
 
+  # TODO(phawkins): add unit_diagonal support to triangular_solve, use it here
+  # instead of explicit masking of l.
+  l = np.tril(lu, -1)[..., :, :m] + np.eye(m, m, dtype=dtype)
   x = lax_linalg.triangular_solve(l, x, left_side=True, lower=True)
   x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=False)
 
-  return x[..., 0] if a_ndims != b_ndims else x
+  return x[..., 0] if a_ndims == b_ndims + 1 else x
 
 
 for func in get_module_functions(onp.linalg):
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
index c907ee53a..e6cf99c7f 100644
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -365,6 +365,9 @@ MIXED_ADVANCED_INDEXING_TESTS = [
                                          None,
                                          onp.array([-1, 2]))),
      ]),
+    ("IntArrayWithInt32Type",
+     [IndexSpec(shape=(3, 4), indexer=(Ellipsis, onp.array(1, dtype=onp.int32)))
+     ]),
 ]
 
 class IndexingTest(jtu.JaxTestCase):
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index e7e7422c7..878153410 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -422,6 +422,8 @@ class NumpyLinalgTest(jtu.JaxTestCase):
           ((1, 1), (1, 1)),
           ((4, 4), (4,)),
           ((8, 8), (8, 4)),
+          ((1, 2, 2), (3, 2)),
+          ((2, 1, 3, 3), (2, 4, 3, 4)),
       ]
       for dtype in float_types() + complex_types()
       for rng in [jtu.rand_default()]))