backfill eager_sharding in set_metadata

PiperOrigin-RevId: 832407389

Author: Cristian Garcia

flax/nnx/variablelib.py

@@ -1261,6 +1261,13 @@ def set_metadata(self, *args, **kwargs) -> None:
           f'Cannot change `is_mutable` metadata, expected {self.is_mutable}, '
           f'got {metadata["is_mutable"]}'
         )
+      if 'eager_sharding' not in metadata:
+        metadata['eager_sharding'] = self.eager_sharding
+      if metadata['eager_sharding'] != self.eager_sharding:
+        raise ValueError(
+          f'Cannot change `eager_sharding` metadata, expected '
+          f'{self.eager_sharding}, got {metadata["eager_sharding"]}'
+        )
       self._var_metadata = metadata
     elif len(args) == 2:
       name, value = args