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

Author: Cristian Garcia

src/MaxText/layers/initializers.py

@@ -64,9 +64,9 @@ def variable_to_logically_partitioned(variable: nnx.VariableState):
   """Wraps an NNX variable's value in `nn.LogicallyPartitioned`.
 
   This function inspects the metadata of an `nnx.VariableState` object. If
-  sharding information ('sharding' or 'sharding_names') is present, it wraps
-  the variable's value in `nn.LogicallyPartitioned` to apply the specified
-  sharding constraints.
+  sharding information ('out_sharding', 'sharding' or 'sharding_names') is
+  present, it wraps the variable's value in `nn.LogicallyPartitioned` to apply
+  the specified sharding constraints.
 
   It handles special cases for `aqt_tensor.QTensor` and variables of type
   `_overwrite_with_gradient` by returning their values directly without
@@ -85,14 +85,18 @@ def variable_to_logically_partitioned(variable: nnx.VariableState):
     return variable.value
 
   metadata = variable.get_metadata()
-  if "sharding" in metadata or "sharding_names" in metadata:
-    if "sharding_names" in metadata:
-      sharding_names = metadata["sharding_names"]
-    else:
-      sharding_names = metadata["sharding"]
+  out_sharding = None
+  if "out_sharding" in metadata:
+    out_sharding = metadata["out_sharding"]
+  elif "sharding_names" in metadata:
+    out_sharding = metadata["sharding_names"]
+  elif "sharding" in metadata:
+    out_sharding = metadata["sharding"]
+
+  if out_sharding is not None:
     return nn.LogicallyPartitioned(  # type: ignore[wrong-keyword-args]
         variable.value,
-        sharding_names,  # type: ignore[arg-type]
+        out_sharding,  # type: ignore[arg-type]
         mesh=metadata.get("mesh"),
         rules=metadata.get("rules"),
     )