jax

6b253b2f - [shardy] Fix cases in shardy where you have a nullary function with partially specified out_shardings (i.e. some out_sharding's are None and others are NamedShardings).

[shardy] Fix cases in shardy where you have a nullary function with partially specified out_shardings (i.e. some out_sharding's are None and others are NamedShardings). In this case, the returned out_shardings should all be NamedSharding (because of NamedSharding's presence in some out_sharding's). PiperOrigin-RevId: 714681941