Merge pull request #4983 from google:propagate-eager-sharding

PiperOrigin-RevId: 811550398

Author: Flax Authors

flax/nnx/variablelib.py

@@ -285,7 +285,7 @@ def __init__(
     # shard the value if applicable
     do_eager_sharding = config.flax_always_shard_variable
     if 'eager_sharding' in metadata:
-      do_eager_sharding = metadata.pop('eager_sharding')
+      do_eager_sharding = metadata['eager_sharding']
     if do_eager_sharding and 'sharding_names' in metadata:
       value = core_spmd.shard_value(
         value, metadata['sharding_names'], metadata.get('sharding_rules', None),