fix: Fix memory scrambling bug in reshape_kernel hooks for scanned layers and MoE experts

Author: khatwanimohit

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -505,10 +505,9 @@ def pad_hf_embedding_layer(input_tensor, target_shape):
 
   def reshape_kernel(input_tensor, target_shape):
     if saving_to_hf:
-      flipped_target_shape = np.flip(np.array(target_shape))
-      return input_tensor.reshape(flipped_target_shape).T
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
     else:
-      return input_tensor.T.reshape(target_shape)
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
 
   def scale_rmsnorm_layer(input_tensor, target_shape):
     if saving_to_hf:
@@ -773,10 +772,9 @@ def pad_embedding_layer(input_tensor, target_shape):
   def reshape_kernel(input_tensor, target_shape):
     """Reshapes and transposes kernel weights between MaxText and HF."""
     if saving_to_hf:
-      flipped_target_shape = np.flip(np.array(target_shape))
-      return input_tensor.reshape(flipped_target_shape).T
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
     else:
-      return input_tensor.T.reshape(target_shape)
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
 
   def reshape_bias(input_tensor, target_shape=None):
     """Reshapes biases between MaxText 2D (heads, dim) and HF 1D (hidden)."""
@@ -1019,10 +1017,9 @@ def transpose(input_tensor, target_shape=None):
 
   def reshape_kernel(input_tensor, target_shape):
     if saving_to_hf:
-      flipped_target_shape = np.flip(np.array(target_shape))
-      return input_tensor.reshape(flipped_target_shape).T
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
     else:
-      return input_tensor.T.reshape(target_shape)
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
 
   def permute_conv(input_tensor, target_shape=None):
     # MT: [K, 1, C] <-> HF: [C, 1, K]
@@ -1174,10 +1171,9 @@ def DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=Fal
   def reshape_kernel(input_tensor, target_shape):
     """Reshapes and transposes kernel weights between MaxText and HF."""
     if saving_to_hf:
-      flipped_target_shape = np.flip(np.array(target_shape))
-      return input_tensor.reshape(flipped_target_shape).T
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
     else:
-      return input_tensor.T.reshape(target_shape)
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
 
   num_main_layers = config["num_hidden_layers"]
   first_num_dense_layers = config["first_k_dense_replace"]
@@ -1362,10 +1358,9 @@ def transpose(input_tensor, target_shape=None):
   def reshape_kernel(input_tensor, target_shape):
     """Reshapes and transposes kernel weights between MaxText and HF."""
     if saving_to_hf:
-      flipped_target_shape = np.flip(np.array(target_shape))
-      return input_tensor.reshape(flipped_target_shape).T
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
     else:
-      return input_tensor.T.reshape(target_shape)
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
 
   def reshape_bias(input_tensor, target_shape=None):
     """Reshapes biases between MaxText 2D (heads, dim) and HF 1D (hidden)."""
@@ -1971,10 +1966,9 @@ def adjust_rope(input_tensor, target_shape):
 
   def reshape_kernel(input_tensor, target_shape):
     if saving_to_hf:
-      flipped_target_shape = np.flip(np.array(target_shape))
-      return input_tensor.reshape(flipped_target_shape).transpose()
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
     else:
-      return input_tensor.transpose().reshape(target_shape)
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
 
   # caveat: hook order does affect result
   # to_huggingface
@@ -2549,10 +2543,9 @@ def pad_hf_embedding_layer(input_tensor, target_shape):
 
   def reshape_kernel(input_tensor, target_shape):
     if saving_to_hf:
-      flipped_target_shape = np.flip(np.array(target_shape))
-      return input_tensor.reshape(flipped_target_shape).T
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
     else:
-      return input_tensor.T.reshape(target_shape)
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
 
   def scale_rmsnorm_layer(input_tensor, target_shape):
     # Shift of 1.0 is now folded into Gemma 4 text and vision checkpoint weights
@@ -2801,10 +2794,9 @@ def OLMO3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False,
   # Standard Transpose for Kernels (HF: [Out, In] <-> MaxText: [In, Out])
   def reshape_kernel(input_tensor, target_shape):
     if saving_to_hf:
-      flipped_target_shape = np.flip(np.array(target_shape))
-      return input_tensor.reshape(flipped_target_shape).T
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
     else:
-      return input_tensor.T.reshape(target_shape)
+      return np.swapaxes(input_tensor, -1, -2).reshape(target_shape)
 
   # Identity mapping for Norms
   # Olmo3 checkpoints typically have weights ~1.0.