Problem with vmapping rootfinder

Hello Everyone.

First of all, thanks for this great library!

I'm not sure if the following issue is with jaxopt or jax itself, but I started having problems applying vmap to a root solver.
I tried following the example in 'gradient_descent_test.py' (particularly 'test_jit_and_vmap').
However, I'm getting a TracerArrayConversionError when I try to evaluate the vmapped function.

```
TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Tracedwith
with val = DeviceArray([[-17.14141909],[ 58.908974 ],....]], dtype=float64)
batch_dim = 0
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
```

Below is a simple example to reproduce the issue. Any help would be greatly appreciated.
Thanks in advance.
John

```
import jax
from jax.config import config; config.update("jax_enable_x64", True)
import jax.numpy as np
from jax import random,vmap

from jaxopt import linear_solve
from jaxopt import ScipyRootFinding

def func(x, params):
a,b = params
return (x - a) * (x - b)
kwargs = {'implicit_diff_solve':linear_solve.solve_normal_cg, 'method':'hybr', 'tol':1e-10}
rootfinder = ScipyRootFinding(optimality_fun=func, **kwargs)

init = np.zeros(1)
def solve(params):
root, info = rootfinder.run(init, params)
return root
key = random.PRNGKey(1235711)
param_list = random.uniform(key, (100, 2), minval = -10, maxval = 10)

root = solve(param_list[0]) # this is ok
root_list = vmap(solve)(param_list) # this doesn't work
```

p.s. I'm using the following packages
python 3.8.8
jax 0.2.19
jaxlib 0.1.70
jaxopt 0.1