`NameError: unbound axis name` when running example in documentation

I'm running the notebook `custom_loop_pmap_example.ipynb` from this [example](https://jaxopt.github.io/stable/notebooks/distributed/custom_loop_pmap_example.html) in the documentation.

When option `use_pmap=True` it produces the following error:

```
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[39], line 9
6 print("linesearch (ignored if `stepsize` > 0):", LINESEARCH)
7 print()
----> 9 errors, step_times, compile_time = run()
10 print('Average speed-up (ignoring compile):',
11 round((step_times['without_pmap'] [/](https://vscode-remote+ssh-002dremote-002bwestbun.vscode-resource.vscode-cdn.net/) step_times['with_pmap']).mean(), 2))

Cell In[38], line 15, in run()
13 exp_name: str = f"{'with' if use_pmap else 'without'}_pmap"
14 print(exp_name)
---> 15 _errors, _step_times, _compile_time = fit(data=data,
16 init_params=init_params,
17 stepsize=STEPSIZE,
18 linesearch=LINESEARCH,
19 use_pmap=use_pmap)
20 errors[exp_name] = _errors
21 step_times[exp_name] = _step_times

Cell In[37], line 49, in fit(data, init_params, stepsize, linesearch, use_pmap)
46 update = jax.jit(solver.update)
48 # Initialize solver state.
---> 49 state = solver.init_state(init_params, data=data)
50 params = init_params
51 # If using `pmap` for data-parallel training, model parameters are typically
52 # replicated across devices.

File [~/jax/lib/python3.10/site-packages/jaxopt/_src/lbfgs.py:284](https://vscode-remote+ssh-002dremote-002bwestbun.vscode-resource.vscode-cdn.net/home/ubuntu/workplace/learn/jaxopt/docs/notebooks/distributed/~/jax/lib/python3.10/site-packages/jaxopt/_src/lbfgs.py:284), in LBFGS.init_state(self, init_params, *args, **kwargs)
275 dtype = tree_single_dtype(init_params)
276 state_kwargs = dict(
277 s_history=init_history(init_params, self.history_size),
278 y_history=init_history(init_params, self.history_size),
(...)
282 stepsize=jnp.asarray(self.max_stepsize, dtype=dtype),
283 )
--> 284 (value, aux), grad = self._value_and_grad_with_aux(init_params, *args, **kwargs)
285 return LbfgsState(value=value,
286 grad=grad,
287 error=jnp.asarray(jnp.inf),
288 **state_kwargs,
289 aux=aux,
290 failed_linesearch=jnp.asarray(False))

File [~/jax/lib/python3.10/site-packages/jaxopt/_src/base.py:62](https://vscode-remote+ssh-002dremote-002bwestbun.vscode-resource.vscode-cdn.net/home/ubuntu/workplace/learn/jaxopt/docs/notebooks/distributed/~/jax/lib/python3.10/site-packages/jaxopt/_src/base.py:62), in _add_aux_to_value_and_grad..value_and_grad_with_aux(*a, **kw)
61 def value_and_grad_with_aux(*a, **kw):
---> 62 v, g = value_and_grad(*a, **kw)
63 return (v, None), g

Cell In[35], line 6, in pmean..wrapper(*args, **kwargs)
4 @functools.wraps(fun)
5 def wrapper(*args, **kwargs):
----> 6 return jax.tree_map(maybe_pmean, fun(*args, **kwargs))

File [~/jax/lib/python3.10/site-packages/jax/_src/tree_util.py:210](https://vscode-remote+ssh-002dremote-002bwestbun.vscode-resource.vscode-cdn.net/home/ubuntu/workplace/learn/jaxopt/docs/notebooks/distributed/~/jax/lib/python3.10/site-packages/jax/_src/tree_util.py:210), in tree_map(f, tree, is_leaf, *rest)
...
-> 2530 raise NameError(
2531 f'unbound axis name: {axis_name}. The following axis names (e.g. defined '
2532 f'by pmap) are available to collective operations: {named_axes}')

NameError: unbound axis name: b. The following axis names (e.g. defined by pmap) are available to collective operations: []
```

This are the versions of jax and jaxopt in my system:

```
jax 0.4.11
jaxlib 0.4.7+cuda11.cudnn86
jaxopt 0.7
```