jax

9bf9b863 - [pmap] Update migrate pmap docs to include sharded reduction issues.