jax solver failing tests

[martinjrobins]

Describe the bug

Some tests in TestJaxBDFSolver are now failing.

File "/home/runner/work/PyBaMM/PyBaMM/pybamm/solvers/jax_bdf_solver.py", line 169, in _bdf_init
    state['newton_tol'] = jnp.max((10 * EPS / rtol, jnp.min((0.03, rtol ** 0.5))))
  File "/home/runner/work/PyBaMM/PyBaMM/.tox/tests/lib/python3.7/site-packages/jax/numpy/lax_numpy.py", line 1670, in reduction
    _check_arraylike(name, a)
  File "/home/runner/work/PyBaMM/PyBaMM/.tox/tests/lib/python3.7/site-packages/jax/numpy/lax_numpy.py", line 297, in _check_arraylike
    raise TypeError(msg.format(fun_name, type(arg), pos))
jax.traceback_util.FilteredStackTrace: TypeError: min requires ndarray or scalar arguments, got <class 'tuple'> at position 0.

From @tlestang: It looks like jnp.min() is given a tuple when its expecting a array. Is this maybe a change in the latest jax @martinjrobins ?

To Reproduce

Run tests with recent version of jax

Expected behaviour

tests pass