flax

de47d9b7 - Merge pull request #5070 from google:variable-explicit-sharding-support