flax
Add support for `jax.sharding.PartitionSpec.UNCONSTRAINED` in logical specification
#3902
Merged
Merged