Repository: google/jax

Issue:
Incorrect batching rule with custom_linear_solve when has_aux=True?
### Description

Since [this commit](https://github.com/google/jax/commit/73ed511d3945246467b0e3e12da326d40a3fe153), a bunch of tests that used to pass in JAXopt are now failing with `ValueError: too many values to unpack`. The issue seems to come from the fact that `custom_linear_solve` has an incorrect batching rule when `has_aux=True`. Here's a minimal repro:
```python
import numpy as np
import jax.numpy as jnp
import jax
from jax.scipy.sparse.linalg import cg

rng = np.random.RandomState(0)

M = rng.randn(5, 5)
A = np.dot(M, M.T)
matvec = lambda x: (jnp.dot(A, x[0]), jnp.dot(A, x[1]))

def f(b):
  return cg(matvec, (b, b))[0]

b = rng.randn(5)
jax.jacrev(f)(b)
```
Error obtained:
```
 Traceback (most recent call last):
  File "/Users/mblondel/Desktop/playground/cg_bug.py", line 16, in <module>
    jax.jacrev(f)(b)
  File "/Users/mblondel/Desktop/playground/cg_bug.py", line 13, in f
    return cg(matvec, (b, b))[0]
  File "/Users/mblondel/Desktop/projects/jax/jax/_src/scipy/sparse/linalg.py", line 290, in cg
    return _isolve(_cg_solve,
  File "/Users/mblondel/Desktop/projects/jax/jax/_src/scipy/sparse/linalg.py", line 229, in _isolve
    x_maybe_info = lax.custom_linear_solve(
jax._src.source_info_util.JaxStackTraceBeforeTransformation: ValueError: too many values to unpack (expected 1)

[...]

  File "/Users/mblondel/Desktop/projects/jax/jax/interpreters/batching.py", line 710, in _match_axes_jaxpr
    out_axis_dest, = out_axes_dest
jax._src.traceback_util.UnfilteredStackTrace: ValueError: too many values to unpack (expected 1)

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/mblondel/Desktop/playground/cg_bug.py", line 16, in <module>
    jax.jacrev(f)(b)
  File "/Users/mblondel/Desktop/projects/jax/jax/_src/api.py", line 1324, in jacfun
    jac = vmap(pullback)(_std_basis(y))
ValueError: too many values to unpack (expected 1)
```

If I come back to before the aforementioned commit, this script works, since `cg` didn't use `custom_linear_solve(..., has_aux=True)`.

Also note that I use `(b, b)` to define a pytree. When the input is just an array, there is no error.

### What jax/jaxlib version are you using?

main branch

### Which accelerator(s) are you using?

CPU

### Additional system info

3.8, Mac

### NVIDIA GPU info

_No response_


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/control_flow/solves.py b/jax/_src/lax/control_flow/solves.py
--- a/jax/_src/lax/control_flow/solves.py
+++ b/jax/_src/lax/control_flow/solves.py
@@ -387,7 +387,9 @@ def _linear_solve_batching_rule(spmd_axis_name, axis_size, axis_name, main_type,
   (matvec, vecmat, solve, solve_t) = jaxprs
   (matvec_bat, vecmat_bat, solve_bat, solve_t_bat) = params_bat
 
-  num_aux = len(solve.out_avals) - len(matvec.out_avals)
+  # number of operator out avals is assumed to be the same for matvec/vecmat
+  num_operator_out_avals = len(matvec.out_avals)
+  num_aux = len(solve.out_avals) - num_operator_out_avals
   # Fixpoint computation of which parts of x and b are batched; we need to
   # ensure this is consistent between all four jaxprs
   b_bat = orig_b_bat
@@ -402,21 +404,23 @@ def _linear_solve_batching_rule(spmd_axis_name, axis_size, axis_name, main_type,
       x_bat_out = solve_x_bat
     else:
       vecmat_jaxpr_batched, vecmat_x_bat = batching.batch_jaxpr(
-          vecmat, axis_size, vecmat_bat + b_bat, instantiate=x_bat,
+          vecmat, axis_size, vecmat_bat + b_bat, instantiate=b_bat,
           axis_name=axis_name, spmd_axis_name=spmd_axis_name, main_type=main_type)
       # batch all aux data by default
       x_bat_out = _map(operator.or_, vecmat_x_bat + [True] * num_aux, solve_x_bat)
+    # keep a slice of only the linear operator part of solve's avals
+    x_bat_noaux = x_bat_out[:num_operator_out_avals]
 
     # Apply matvec and solve_t -> new batched parts of b
     matvec_jaxpr_batched, matvec_b_bat = batching.batch_jaxpr(
-        matvec, axis_size, matvec_bat + x_bat_out, instantiate=b_bat,
+        matvec, axis_size, matvec_bat + x_bat_noaux, instantiate=b_bat,
         axis_name=axis_name, spmd_axis_name=spmd_axis_name, main_type=main_type)
     if solve_t is None:
       solve_t_jaxpr_batched = None
       b_bat_out = _map(operator.or_, matvec_b_bat, orig_b_bat)
     else:
       solve_t_jaxpr_batched, solve_t_b_aux_bat = batching.batch_jaxpr(
-          solve_t, axis_size, solve_t_bat + x_bat_out, instantiate=b_bat,
+          solve_t, axis_size, solve_t_bat + x_bat_noaux, instantiate=x_bat_out,
           axis_name=axis_name, spmd_axis_name=spmd_axis_name, main_type=main_type)
       assert len(solve_t_b_aux_bat) == len(orig_b_bat) + num_aux
       solve_t_b_bat, _ = split_list(solve_t_b_aux_bat, [len(orig_b_bat)])


Test patch:
diff --git a/tests/custom_linear_solve_test.py b/tests/custom_linear_solve_test.py
--- a/tests/custom_linear_solve_test.py
+++ b/tests/custom_linear_solve_test.py
@@ -467,6 +467,23 @@ def linear_solve(a, b):
       return matrix_free_solve(partial(high_precision_dot, a), b)
     jtu.check_grads(linear_solve, (a, b), order=1, rtol=3e-3, modes=['rev'])
 
+  def test_custom_linear_solve_batching_with_aux(self):
+    def solve(mv, b):
+      aux = (np.array(1.), True, 0)
+      return mv(b), aux
+
+    def solve_aux(x):
+      matvec = lambda y: tree_util.tree_map(partial(jnp.dot, A), y)
+      return lax.custom_linear_solve(matvec, (x, x), solve, solve, symmetric=True, has_aux=True)
+
+    rng = self.rng()
+    A = rng.randn(3, 3)
+    A = A + A.T
+    b = rng.randn(3, 3)
+
+    # doesn't crash
+    jax.vmap(solve_aux)(b)
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())