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.


Discussion:
I also observed this problem and solved it by making `A` and `b` have the same batch_shape (using `np.broadcast_to`).
@fehiepsi Hmm that's not actually working for me. I now have shapes `(3, 2, 2)` and `(3, 2)` for A and b, respectively but I'm still seeing the same error.

And `A.shape == (3, 2, 2)` `b.shape == (3, 2, 1)` doesn't work either.
Sorry, I used `jax.scipy.linalg.solve_triangular` so the behaviour is a bit difference. The case `(3, 2, 2)` and `(3, 2)`/`(3, 2, 1)` does not work with `jax.numpy.solve` in my system too.
@fehiepsi Ah, that would explain it!
Yes, `jax.numpy.linalg.solve` simply doesn't support broadcasting at the moment.

As a workaround, you could most likely wrap a `vmap` around an unbatched solve; I see no reason why that wouldn't work. The underlying primitives should support batching, it's only the Python numpy layer on top that does not support batching.
@hawkinsp I can see how `vmap` would work when batching over a single dimension, but is there a solution for batching across multiple dimensions? For example `A.shape == (5, 3, 2, 2) b.shape == (5, 3, 2, 1)`.
Well, hopefully it's moot now because `np.linalg.solve` should support broadcasting.

But you could use `vmap` with multiple batch dimensions, either by flattening them with a `reshape`, or by applying `vmap` repeatedly. If this is a common pattern, we might want to add support to `vmap` itself to map over multiple dimensions.
@hawkinsp Roger that! Yeah, I would totally vote for vmap supporting batching over a variable number of dimensions, ie. "batch over all but last n dimensions".
Hi!
I'm not sure if it is better to open a new issue or write here (in case just let me know) but I still seem to have problem with solve and the broadcasting semantics.

In my problem I have to compute the bilinear `y' S y` when the vector `y` has shape `(dim, )` and `S` is a batch of matrices of shape `(comp, dim, dim)`.
My code works with the standard numpy's implementation of `linalg.solve` but not with jax's:

```python
import numpy as np
import jax.numpy as jnp

from jax import random
RNG = random.PRNGKey(42)
keyY, keyS = random.split(RNG)

dim = 3; comp = 4

y = jnp.expand_dims(random.uniform(keyY, shape=(dim, )), 0)
S = random.uniform(keyS, shape=(comp, dim, dim))

print(y.shape)
print(S.shape)

for solver in [np.linalg.solve, jnp.linalg.solve]:
    print(jnp.sum(y * solver(S, y), axis=1).shape)
```

This works for pure numpy code, but for jax it fails, and from what I can see (output trace provided below) it appears that `jnp.linalg.solve` is not performing the broadcasting over the first dimension of y. What am I doing wrong?

I am working with jax version `0.2.7` and numpy `1.19.2`.

Error Trace:
```
---------------------------------------------------------------------------
FilteredStackTrace                        Traceback (most recent call last)
<ipython-input-5-766e89175182> in <module>
      1 for solver in [np.linalg.solve, jnp.linalg.solve]:
----> 2     print(jnp.sum(y * solver(S, y), axis=1).shape)
      3 

~/.miniconda3/envs/pyintel/lib/python3.7/site-packages/jax/_src/numpy/linalg.py in solve(a, b)
    449   a, b = _promote_arg_dtypes(jnp.asarray(a), jnp.asarray(b))
--> 450   return lax_linalg._solve(a, b)
    451 

~/.miniconda3/envs/pyintel/lib/python3.7/site-packages/jax/_src/lax/linalg.py in _solve(a, b)
    265     # b.shape == [..., m, k]
--> 266     return api.vmap(custom_solve, b.ndim - 1, max(a.ndim, b.ndim) - 1)(b)
    267 

~/.miniconda3/envs/pyintel/lib/python3.7/site-packages/jax/_src/lax/control_flow.py in custom_linear_solve(matvec, b, solve, transpose_solve, symmetric)                                                                  
   2179   matvec_jaxpr, matvec_consts, out_tree = _initial_style_jaxpr(
-> 2180       _shape_checked(matvec, "matvec"), in_args_tree, b_avals)
   2181   _check_tree("matvec", "b", out_tree, tree)

~/.miniconda3/envs/pyintel/lib/python3.7/site-packages/jax/_src/lax/control_flow.py in _initial_style_jaxpr(fun, in_tree, in_avals)                                                                                       
     71 def _initial_style_jaxpr(fun: Callable, in_tree, in_avals):
---> 72   jaxpr, out_avals, consts, out_tree = _initial_style_open_jaxpr(fun, in_tree, in_avals)
     73   closed_jaxpr = core.ClosedJaxpr(pe.convert_constvars_jaxpr(jaxpr), ())

~/.miniconda3/envs/pyintel/lib/python3.7/site-packages/jax/_src/lax/control_flow.py in _initial_style_open_jaxpr(fun, in_tree, in_avals)
     66   wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
---> 67   jaxpr, out_avals, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals)
     68   return jaxpr, out_avals, consts, out_tree()

~/.miniconda3/envs/pyintel/lib/python3.7/site-packages/jax/_src/lax/control_flow.py in f(x)
   2174       y = fun(x)
-> 2175       _check_shapes(name, "b", y, b_flat)
   2176       return y

~/.miniconda3/envs/pyintel/lib/python3.7/site-packages/jax/_src/lax/control_flow.py in _check_shapes(func_name, expected_name, actual, expected)                                                                          
   2125     raise ValueError(
-> 2126         f"{func_name}() output shapes must match {expected_name}, "
   2127         f"got {actual_shapes} and {expected_shapes}")

FilteredStackTrace: ValueError: matvec() output shapes must match b, got [(4, 3)] and [(1, 3)]

The stack trace above excludes JAX-internal frames.
```

It looks like this is a bug related to broadcasting within `np.linalg.solve`. Shorter repro:
```python
import numpy as np
import jax.numpy as jnp

S = np.random.rand(4, 3, 3)
y = np.random.rand(1, 3)

print(np.linalg.solve(S, y).shape)
# (4, 3)

print(jnp.linalg.solve(S, y).shape)
# ValueError: matvec() output shapes must match b, got [(4, 3)] and [(1, 3)]
```

Resolving patch:
diff --git a/jax/_src/lax/linalg.py b/jax/_src/lax/linalg.py
--- 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))


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -706,7 +706,7 @@ def tensor_maker():
           ((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))