flax

c05d004d - support explicit sharding in eager sharding