Merge pull request #3148 from AI-Hypercomputer:chengnuojin-fix-debug

PiperOrigin-RevId: 871547910

Author: Google-ML-Automation

src/MaxText/sharding.py

@@ -78,7 +78,7 @@ def maybe_shard_with_logical(
   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)
+    log_key = (str(jax.typeof(inputs)), tuple(logical_axes), extra_stack_level)
 
     if log_key not in _LOGGED_LOGICAL_AXES:
       max_logging.info(f"Logical:  {log_key[0]:.<60} {log_key[1]}", stacklevel=3 + extra_stack_level)