Update Pmap_Cookbook.ipynb
Update Pmap_Cookbook.ipynb

I am doing this mostly to trigger another run of CI; there was a problem with our CI runs and this PR appears to be failing tests.
google/jax,google/jax,google/jax,google/jax