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.
diff --git a/jax/_src/lax/linalg.py b/jax/_src/lax/linalg.py
index 49c4b9d5e..ac7ba3f5f 100644
--- a/jax/_src/lax/linalg.py
+++ b/jax/_src/lax/linalg.py
@@ -253,6 +253,12 @@ def _check_solve_shapes(a, b):
 def _solve(a, b):
   _check_solve_shapes(a, b)
 
+  # Broadcast leading dimensions of b to the shape of a, as is required by
+  # custom_linear_solve.
+  out_shape = tuple(d_a if d_b == 1 else d_b
+                    for d_a, d_b in zip(a.shape[:-1] + (1,), b.shape))
+  b = jnp.broadcast_to(b, out_shape)
+
   # With custom_linear_solve, we can reuse the same factorization when
   # computing sensitivities. This is considerably faster.
   lu_, _, permutation = lu(lax.stop_gradient(a))
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 6dabfe6f0..4f35fef50 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -706,7 +706,7 @@ class NumpyLinalgTest(jtu.JaxTestCase):
           ((4, 4), (4,)),
           ((8, 8), (8, 4)),
           ((1, 2, 2), (3, 2)),
-          ((2, 1, 3, 3), (2, 4, 3, 4)),
+          ((2, 1, 3, 3), (1, 4, 3, 4)),
           ((1, 0, 0), (1, 0, 2)),
       ]
       for dtype in float_types + complex_types))
['jax/_src/lax/linalg.py', 'tests/linalg_test.py']