flax

[nnx] enable sharding transformation on integer prefixes

#4185

Merged