Pau Pereira

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

2

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...