Tags: Zirunis/flax
Tags
Rename sharding_names to out_sharding in NNX Variable metadata This CL renames the sharding_names attribute to out_sharding for better consistency with the sharding API. The new name more clearly indicates the purpose of this metadata field. ## Changes - Bump Flax version to 0.12.4 - Core changes in variablelib.py: - Add sharding_names to out_sharding metadata remapping for backward compatibility - Add deprecated sharding_names property that returns out_sharding with a warning - Update nnx/spmd.py, core/spmd.py, core/meta.py, linen/spmd.py to use out_sharding - Update all NNX tests to use the new attribute name - Update qwix flax_util.py to check for out_sharding first, with fallback to sharding_names - Update maxtext initializers.py to check for out_sharding first - Update documentation and examples to use out_sharding ## Backward Compatibility Existing code using sharding_names will continue to work via: - Metadata remapping during Variable creation - Deprecated Variable.sharding_names property PiperOrigin-RevId: 869269899
Merge pull request google#5149 from IvyZX:release PiperOrigin-RevId: 846407227
Merge pull request google#5095 from google:jax-minver-0.8.1 PiperOrigin-RevId: 834401927
Merge pull request google#4983 from google:propagate-eager-sharding PiperOrigin-RevId: 811550398
Merge pull request google#4878 from google:ver-0.11.1 PiperOrigin-RevId: 792761136
Merge pull request google#4854 from google:mutable-array-guide PiperOrigin-RevId: 788577118
Merge pull request google#4798 from google:ver-0.10.7 PiperOrigin-RevId: 778346925
Merge pull request google#4724 from google:0.10.6 PiperOrigin-RevId: 750649299
PreviousNext