fix comments

Author: Sharon Yu

src/MaxText/maxtext_utils.py

@@ -1246,21 +1246,18 @@ def print_shardings_params(params, params_sharding, mesh, logical_annotations=No
     params = {"params": params}
   if not hasattr(params_sharding, "params"):
     params_sharding = {"params": params_sharding}
+  if logical_annotations and not hasattr(logical_annotations, "params"):
+    logical_annotations = {"params": logical_annotations}
 
   leaves_params, _ = jax.tree_util.tree_flatten_with_path(params)
   leaves_sharding, _ = jax.tree_util.tree_flatten_with_path(params_sharding)
+  leaves_logical, _ = jax.tree_util.tree_flatten_with_path(logical_annotations.params)
 
-  leaves_logical = []
-  has_logical = False
-  if logical_annotations and hasattr(logical_annotations, "params"):
-    try:
-      leaves_logical, _ = jax.tree_util.tree_flatten_with_path(logical_annotations.params)
-      if len(leaves_params) == len(leaves_logical):
-        has_logical = True
-      else:
-        max_logging.warning("Warning: Logical annotations tree structure mismatch. Skipping logical info.")
-    except Exception as e:  # pylint: disable=broad-exception-caught
-      max_logging.warning(f"Warning: Failed to process logical annotations: {e}. Skipping logical info.")
+  for i, ((path, leaf_val), (_, leaf_sharding)) in enumerate(zip(leaves_params, leaves_sharding)):
+    path_str = "/".join(str(p.key if hasattr(p, "key") else p.name) for p in path)
+    shape = jax.typeof(leaf_val)
+    pspec = sharding.remove_size_one_mesh_axis(leaf_sharding.spec, mesh)
+    pspec_str = str(tuple(pspec))
 
   if not has_logical:
     leaves_logical = [(None, None)] * len(leaves_params)

src/MaxText/model_creation_utils.py

@@ -160,8 +160,8 @@ def create_sharded_state():
       maxtext_utils.print_shardings_params(
           params=sharded_state,
           params_sharding=out_shardings,
-          logical_annotations=specs,
           mesh=model.mesh,
+          logical_annotations=specs,
       )
     if config.load_parameters_path:
       try:

src/MaxText/sharding.py

@@ -31,35 +31,28 @@
 
 
 _LOGGED_ACTIVATION_SHARDINGS = set()
+_LOGGED_LOGICAL_AXES = set()
 
 
 def get_input_data_sharding(config, mesh):
   """Get the input data sharding for the model"""
   return create_sharding(mesh, config.input_data_sharding_logical_axes, rules=config.logical_axis_rules)
 
 
-def maybe_shard_with_name(
-    inputs, named_sharding, shard_mode, debug_sharding=False, extra_stack_level=0, logical_axes=None
-):
+def maybe_shard_with_name(inputs, named_sharding, shard_mode, debug_sharding=False, extra_stack_level=0):
   """
   In auto shardmode, this function hints inputs follow given named_sharding.
   In explicit shardmode, this function enforces inputs following named_sharding.
   """
   if inputs is None:
     return None
-  if debug_sharding and isinstance(inputs, Tracer) and isinstance(named_sharding, NamedSharding):
+  if (
+      debug_sharding and isinstance(inputs, Tracer) and isinstance(named_sharding, NamedSharding)
+  ):  # only print pspec for JitTracer
     pspec = remove_size_one_mesh_axis(getattr(named_sharding, "spec"), getattr(named_sharding, "mesh"))
-    if logical_axes is not None:
-      logical_str = str(logical_axes)
-    else:
-      logical_str = "None"
-    shape_str = str(jax.typeof(inputs))
-    log_key = (shape_str, tuple(pspec), extra_stack_level, logical_str)
-
+    log_key = (str(jax.typeof(inputs)), tuple(pspec), extra_stack_level)
     if log_key not in _LOGGED_ACTIVATION_SHARDINGS:
-      max_logging.info(
-          f"Activation: {logical_str:<40} -> {str(tuple(pspec)):<30} {shape_str}", stacklevel=3 + extra_stack_level
-      )
+      max_logging.info(f"{log_key[0]:.<80} {log_key[1]}.", stacklevel=3 + extra_stack_level)
       _LOGGED_ACTIVATION_SHARDINGS.add(log_key)
   if shard_mode == ShardMode.EXPLICIT:
     return reshard(inputs, named_sharding)
@@ -75,14 +68,26 @@ def maybe_shard_with_logical(
   """
   if inputs is None:
     return None
+
   named_sharding = create_sharding(mesh, logical_axes, rules=rules)
+
+  if debug_sharding and isinstance(inputs, Tracer):
+    log_key = (str(jax.typeof(inputs)), logical_axes, extra_stack_level)
+
+    if log_key not in _LOGGED_LOGICAL_AXES:
+      pspec = remove_size_one_mesh_axis(getattr(named_sharding, "spec"), getattr(named_sharding, "mesh"))
+      pspec_str = str(tuple(pspec)) if pspec else "None"
+
+      max_logging.info(f"Logical:  {log_key[0]:.<60} {log_key[1]}", stacklevel=3 + extra_stack_level)
+      max_logging.info(f"{log_key[0]:.<80} {pspec_str}.", stacklevel=3 + extra_stack_level)
+      _LOGGED_LOGICAL_AXES.add(log_key)
+
   return maybe_shard_with_name(
       inputs,
       named_sharding,
       shard_mode,
-      debug_sharding=debug_sharding,
+      debug_sharding=False,
       extra_stack_level=extra_stack_level + 1,
-      logical_axes=logical_axes,
   )