jax

351187d9 - [sharding_in_types] Add support for nary ops to propagate sharding when 1 input is sharded and all others are replicated.