jax

727fe84c - Enter into correct mesh context in shard_map in `_partial_eval_jaxpr_custom_rule`.