jax

7ff6c399 - Remove callable out_specs support from internal `_shard_map`. It existed of pmap_shmap merge but it doesn't use the callable support anymore