Update

Author: hsuan-lun-chiang

src/maxtext/layers/initializers.py

@@ -96,16 +96,13 @@ def variable_to_logically_partitioned(variable: nnx.VariableState):
   if out_sharding is not None:
     if nnx.PARTITION_NAME in metadata:
       partition_name = metadata[nnx.PARTITION_NAME]
-      # Only nnx.Param variables are typically scanned across the param_scan_axis
       scan_axis = metadata.get("param_scan_axis", 0) if variable.type == nnx.Param else 0
 
-      if isinstance(out_sharding, str):
-        out_sharding = [out_sharding]
-      else:
-        out_sharding = list(out_sharding)
+      sharding_list = [out_sharding] if isinstance(out_sharding, str) else list(out_sharding)
+      if partition_name not in sharding_list:
+        sharding_list.insert(scan_axis, partition_name)
 
-      out_sharding.insert(scan_axis, partition_name)
-      out_sharding = tuple(out_sharding)
+      out_sharding = tuple(sharding_list)
 
     return nn.LogicallyPartitioned(  # type: ignore[wrong-keyword-args]
         variable.value,