flax

a7929b37 - Copybara import of the project:

Copybara import of the project: -- 3e963438da492930e808f632fab99ea2b41667c0 by Sam Anklesaria <sanklesaria@openteams.com>: Rename sharding_names to sharding_metadata PiperOrigin-RevId: 853948564