flax

d3754f63 - Merge pull request #4996 from mohsinm-dev:fix-nnx-jit-static-sharding