param_mapping.py

#  Copyright 2023–2026 Google LLC
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""Parameter mappings and transformation hooks for checkpoint conversion.

This module defines the necessary components to convert model checkpoints between
MaxText and Hugging Face formats for various architectures (e.g., Gemma, Qwen).
It provides two key types of mappings for each model:

1.  **Parameter Name Mappings (`PARAM_MAPPING`)**: Dictionaries that map a MaxText
    parameter key to its corresponding Hugging Face parameter(s). These mappings are
    generated by functions like `GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING`.

    **Key: MaxText parameters, with following forms:**
    - `atomic_mt_key`: A single string representing one MaxText parameter.
    - `composite_mt_key`: A tuple of strings representing multiple MaxText parameters. (e.g., GPT-OSS)

    **Value: corresponding Hugging Face parameters, with following forms:**
    - `unscanned`: A single string.
    - `scanned`: A list of strings, to be stacked along the layer axis.
    - `unscanned with expert stacking`: A list of strings, to be stacked along the expert axis.
    - `scanned with expert stacking`: A nested list of strings, to be stacked along both layer and expert axes.
    Note: Expert stacking only applies a subset of MoE models (e.g., Qwen MoE, DeepSeek, Mixtral),
      but not others (e.g., GPT-OSS).

2.  **Hook Functions (`HOOK_FNS`)**: Dictionaries that map a MaxText parameter
    name to a specific transformation function (a "hook"). These hooks handle
    the actual value conversion, which can include operations like reshaping,
    transposing, scaling, or padding tensors to match the target format's
    requirements. These are generated by functions like
    `GEMMA2_MAXTEXT_TO_HF_PARAM_HOOK_FN`.

The main conversion script uses these mappings to systematically transform each
parameter from the source checkpoint and build the target checkpoint.
"""

import warnings
import numpy as np

import jax
import jax.numpy as jnp


def GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
  """Generates a parameter mapping from MaxText to Hugging Face for Gemma3.

  This function creates a dictionary that maps the parameter names from a
  MaxText Gemma3 checkpoint to their corresponding names in the Hugging Face
  `Gemma3ForCausalLM` format. It handles both the text and vision components
  of the model.

  Args:
    config (dict): The Hugging Face model configuration dictionary, which must
      contain 'text_config' and 'vision_config' sub-dictionaries.
    scan_layers (bool, optional): If True, generates mappings for scanned
      layers, where multiple layers are stacked into a single tensor. If False,
      generates mappings for individual, unscanned layers. Defaults to False.

  Returns:
    dict: A mapping where keys are `atomic_mt_key` (single MaxText parameter names). Values
      are either a single Hugging Face parameter name (unscanned form) or a list of
      Hugging Face parameter names (scanned form) for stacked text layers.
  """
  tcfg = config["text_config"]
  vcfg = config["vision_config"]
  Ndec = tcfg["num_hidden_layers"]
  Nvis = vcfg["num_hidden_layers"]

  # pylint: disable=line-too-long
  mapping = {
      # Embedding & final norm
      "params-token_embedder-embedding": "model.language_model.embed_tokens.weight",
      "params-decoder-decoder_norm-scale": "model.language_model.norm.weight",
      # Vision embed & pos
      "params-vision_encoder-Gemma3VisionEncoderLayer_0-embedding-kernel": "model.vision_tower.vision_model.embeddings.patch_embedding.weight",
      "params-vision_encoder-Gemma3VisionEncoderLayer_0-embedding-bias": "model.vision_tower.vision_model.embeddings.patch_embedding.bias",
      "params-vision_encoder-Gemma3VisionEncoderLayer_0-pos_embedding": "model.vision_tower.vision_model.embeddings.position_embedding.weight",
      "params-vision_encoder-Gemma3VisionEncoderLayer_0-Transformer-encoder_norm-scale": "model.vision_tower.vision_model.post_layernorm.weight",
      "params-vision_encoder-Gemma3VisionEncoderLayer_0-Transformer-encoder_norm-bias": "model.vision_tower.vision_model.post_layernorm.bias",
      # Multi-modal projector
      "params-vision_encoder-VisionEmbedder_0-mm_input_projection-w": "model.multi_modal_projector.mm_input_projection_weight",
      "params-vision_encoder-VisionEmbedder_0-mm_soft_embedding_norm-scale": "model.multi_modal_projector.mm_soft_emb_norm.weight",
  }

  vision_params = [
      ("LayerNorm_0-scale", "layer_norm1.weight"),
      ("LayerNorm_0-bias", "layer_norm1.bias"),
      ("LayerNorm_1-scale", "layer_norm2.weight"),
      ("LayerNorm_1-bias", "layer_norm2.bias"),
      ("MultiHeadDotProductAttention_0-query-kernel", "self_attn.q_proj.weight"),
      ("MultiHeadDotProductAttention_0-query-bias", "self_attn.q_proj.bias"),
      ("MultiHeadDotProductAttention_0-key-kernel", "self_attn.k_proj.weight"),
      ("MultiHeadDotProductAttention_0-key-bias", "self_attn.k_proj.bias"),
      ("MultiHeadDotProductAttention_0-value-kernel", "self_attn.v_proj.weight"),
      ("MultiHeadDotProductAttention_0-value-bias", "self_attn.v_proj.bias"),
      ("MultiHeadDotProductAttention_0-out-kernel", "self_attn.out_proj.weight"),
      ("MultiHeadDotProductAttention_0-out-bias", "self_attn.out_proj.bias"),
      ("MlpBlockViT_0-Dense_0-kernel", "mlp.fc1.weight"),
      ("MlpBlockViT_0-Dense_0-bias", "mlp.fc1.bias"),
      ("MlpBlockViT_0-Dense_1-kernel", "mlp.fc2.weight"),
      ("MlpBlockViT_0-Dense_1-bias", "mlp.fc2.bias"),
  ]

  # Vision layers mapping
  for i in range(Nvis):
    for mx, hf in vision_params:
      key = f"params-vision_encoder-Gemma3VisionEncoderLayer_0-Transformer-encoderblock_{i}-{mx}"
      mapping[key] = f"model.vision_tower.vision_model.encoder.layers.{i}.{hf}"

  # Text decoder mapping
  text_params = [
      ("pre_self_attention_norm-scale", "input_layernorm.weight"),
      ("post_self_attention_norm-scale", "post_attention_layernorm.weight"),
      ("self_attention-query_norm-scale", "self_attn.q_norm.weight"),
      ("self_attention-key_norm-scale", "self_attn.k_norm.weight"),
      ("pre_ffw_norm-scale", "pre_feedforward_layernorm.weight"),
      ("post_ffw_norm-scale", "post_feedforward_layernorm.weight"),
      ("self_attention-query-kernel", "self_attn.q_proj.weight"),
      ("self_attention-key-kernel", "self_attn.k_proj.weight"),
      ("self_attention-value-kernel", "self_attn.v_proj.weight"),
      ("self_attention-out-kernel", "self_attn.o_proj.weight"),
      ("mlp-wi_0-kernel", "mlp.gate_proj.weight"),
      ("mlp-wi_1-kernel", "mlp.up_proj.weight"),
      ("mlp-wo-kernel", "mlp.down_proj.weight"),
  ]

  if scan_layers:
    # Gemma3 repeats a 6-layer attention pattern (5 local + 1 global),
    # scanned as layers_0..layers_5 with leftovers in layers_remainder.
    attention_pattern_length = 6
    num_remaining = Ndec % attention_pattern_length
    num_scanned = Ndec - num_remaining

    # Main scanned blocks: params-decoder-layers-layers_{block_idx}-{param}
    for block_idx in range(attention_pattern_length):
      hf_indices = list(range(block_idx, num_scanned, attention_pattern_length))
      for mx, hf in text_params:
        key = f"params-decoder-layers-layers_{block_idx}-{mx}"
        mapping[key] = [f"model.language_model.layers.{i}.{hf}" for i in hf_indices]

    # Remainder layers (unscanned): params-decoder-layers_remainder-layers_{rem_idx}-{param}
    if num_remaining > 0:
      for rem_idx in range(num_remaining):
        hf_layer_idx = num_scanned + rem_idx
        for mx, hf in text_params:
          key = f"params-decoder-layers_remainder-layers_{rem_idx}-{mx}"
          mapping[key] = f"model.language_model.layers.{hf_layer_idx}.{hf}"
  else:
    for i in range(Ndec):
      for mx, hf in text_params:
        key = f"params-decoder-layers_{i}-{mx}"
        mapping[key] = f"model.language_model.layers.{i}.{hf}"

  return mapping


def GEMMA3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
  """Hook functions for Gemma3 parameter conversion.

  This function provides a dictionary of transformation functions (hooks) for
  converting Gemma3 model parameters between MaxText and Hugging Face formats.
  It handles embedding padding/scaling, RMSNorm scaling, kernel reshaping, and
  vision-specific tensor manipulations.

  Args:
    config (dict): The Hugging Face model configuration dictionary.
    scan_layers (bool, optional): Whether the model uses scanned layers.
      Defaults to False.
    saving_to_hf (bool, optional): The direction of conversion. True for
      MaxText to Hugging Face, False for the reverse. Defaults to False.

  Returns:
    dict: A dictionary mapping MaxText parameter names to their corresponding
      transformation functions.
  """
  hooks = {}

  # ---- Embedding pad & scale ----
  def pad_and_scale_embedding(input_tensor, target_shape):
    source_vocab_size, _ = input_tensor.shape
    target_vocab_size, target_hidden_size = target_shape

    # MaxText embedding = original_embedding * sqrt(hidden_size)
    # HF embedding = original_embedding (HF model forward pass applies scaling)
    # Note: config["hidden_size"] is the HF hidden size from the HF config object
    normalizer = np.dtype("bfloat16").type(config["text_config"]["hidden_size"] ** 0.5)

    # Apply scaling first
    if saving_to_hf:  # MaxText to HF
      scaled_tensor = (input_tensor / normalizer).astype(input_tensor.dtype)
    else:  # HF to MaxText
      scaled_tensor = (input_tensor * normalizer).astype(input_tensor.dtype)

    # Handle padding/truncation
    if source_vocab_size > target_vocab_size:
      warnings.warn(
          f"source vocab={source_vocab_size} > target vocab={target_vocab_size}, truncate output layer for MaxText."
      )
      output_tensor = scaled_tensor[:target_vocab_size, :]
    elif source_vocab_size < target_vocab_size:
      warnings.warn(f"source vocab={source_vocab_size} < target vocab={target_vocab_size}, pad output layer for MaxText.")
      padding_shape = (target_vocab_size - source_vocab_size, target_hidden_size)
      # Use jnp.zeros for JAX arrays, np.zeros for numpy arrays
      padding = (
          jnp.zeros(padding_shape, dtype=scaled_tensor.dtype)
          if isinstance(scaled_tensor, jax.Array)
          else np.zeros(padding_shape, dtype=scaled_tensor.dtype)
      )
      output_tensor = (
          jnp.concatenate([scaled_tensor, padding], axis=0)
          if isinstance(scaled_tensor, jax.Array)
          else np.concatenate([scaled_tensor, padding], axis=0)
      )
    else:  # Vocab sizes match
      output_tensor = scaled_tensor

    return output_tensor

  # ---- RMSNorm scale ----
  def scale_rmsnorm(x, target_shape):
    # MaxText norm = HF norm +1; HF norm = MaxText norm -1
    if saving_to_hf:
      return (x - 1.0).reshape(target_shape)
    return (x + 1.0).reshape(target_shape)

  # ---- Generic reshape ----
  def reshape_kernel(x, target_shape):
    if saving_to_hf:
      flipped = np.flip(np.array(target_shape))
      return x.reshape(flipped).T
    else:
      return x.T.reshape(target_shape)

  # ---- Vision reshape ----
  def vis_bias(x, target_shape):
    if saving_to_hf:
      return x.flatten()
    else:
      return x.reshape(target_shape)

  def vision_patch(x, target_shape):
    if saving_to_hf:
      return x.transpose(3, 2, 0, 1)
    else:
      return x.transpose(2, 3, 1, 0)

  def pos_embed(x, target_shape):
    if saving_to_hf:
      return x.squeeze(0)
    return x[None, :, :]

  # ---Embedding & final norm---
  hooks["params-token_embedder-embedding"] = pad_and_scale_embedding
  hooks["params-decoder-decoder_norm-scale"] = scale_rmsnorm
  # [1, 4096, 1152]
  hooks["params-vision_encoder-Gemma3VisionEncoderLayer_0-embedding-kernel"] = vision_patch
  hooks["params-vision_encoder-Gemma3VisionEncoderLayer_0-pos_embedding"] = pos_embed

  hooks["params-vision_encoder-VisionEmbedder_0-mm_input_projection-w"] = lambda x, _: x
  hooks["params-vision_encoder-VisionEmbedder_0-mm_soft_embedding_norm-scale"] = scale_rmsnorm

  # Text layers
  tc = config.get("text_config", {})
  nlayers = tc.get("num_hidden_layers", 0)
  if scan_layers:
    attention_pattern_length = 6
    num_remaining = nlayers % attention_pattern_length
    # Scanned sub-layer prefixes
    prefixes = [f"params-decoder-layers-layers_{block_idx}-" for block_idx in range(attention_pattern_length)]
    # Remainder sub-layer prefixes
    if num_remaining > 0:
      prefixes += [f"params-decoder-layers_remainder-layers_{rem_idx}-" for rem_idx in range(num_remaining)]
  else:
    prefixes = [f"params-decoder-layers_{i}-" for i in range(nlayers)]
  for pref in prefixes:
    # Attention Q/K/V/O
    hooks[pref + "self_attention-query-kernel"] = reshape_kernel
    hooks[pref + "self_attention-key-kernel"] = reshape_kernel
    hooks[pref + "self_attention-value-kernel"] = reshape_kernel
    hooks[pref + "self_attention-out-kernel"] = reshape_kernel
    # Norm scales
    for nm in [
        "pre_self_attention_norm-scale",
        "post_self_attention_norm-scale",
        "self_attention-query_norm-scale",
        "self_attention-key_norm-scale",
        "pre_ffw_norm-scale",
        "post_ffw_norm-scale",
    ]:
      hooks[pref + nm] = scale_rmsnorm
    # MLP
    hooks[pref + "mlp-wi_0-kernel"] = reshape_kernel
    hooks[pref + "mlp-wi_1-kernel"] = reshape_kernel
    hooks[pref + "mlp-wo-kernel"] = reshape_kernel

  # Vision layers
  vc = config.get("vision_config", {})
  nvis = vc.get("num_hidden_layers", 0)
  for i in range(nvis):
    base = f"params-vision_encoder-Gemma3VisionEncoderLayer_0-Transformer-encoderblock_{i}-"
    # Attention kernels & biases
    for qkv in ["query", "key", "value"]:
      hooks[base + f"MultiHeadDotProductAttention_0-{qkv}-kernel"] = reshape_kernel
      hooks[base + f"MultiHeadDotProductAttention_0-{qkv}-bias"] = vis_bias
    # [1152, 1152] -> [16, 72, 1152]
    hooks[base + "MultiHeadDotProductAttention_0-out-kernel"] = reshape_kernel
    hooks[base + "MultiHeadDotProductAttention_0-out-bias"] = vis_bias
    # MLP ViT kernels & biases
    for dense in ["Dense_0", "Dense_1"]:
      hooks[base + f"MlpBlockViT_0-{dense}-kernel"] = reshape_kernel

  return hooks


def GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
  """Returns mapping between MaxText and HuggingFace Gemma2 weight paths.

  Args:
      config (dict): Model configuration dictionary containing at least
        'num_hidden_layers'.
      scan_layers (bool, optional): Whether the MaxText model uses layer
        scanning optimization. When True, decoder layers are stacked into a
        single tensor. Defaults to False.

  Returns:
      dict: A mapping where keys are `atomic_mt_key` (single MaxText parameter name).
        Values are either a single string (unscanned form) or a list of strings
        (scanned form) for stacked layers when `scan_layers=True`.

  Notes:
      - MaxText uses a paired layer approach where two HF decoder layers are
        treated as one MaxText decoder layer.
      - MaxText layer `i` corresponds to HF layers `2i` and `2i+1`.
      - Local components map to even-numbered HF decoder layers (0, 2, 4...).
      - Global components map to odd-numbered HF decoder layers (1, 3, 5...).
  """

  nlayers = config["num_hidden_layers"]
  mapping = {
      "params-token_embedder-embedding": "model.embed_tokens.weight",
      "params-decoder-decoder_norm-scale": "model.norm.weight",
  }
  if scan_layers:
    mapping = {
        **mapping,
        "params-decoder-layers-pre_self_attention_norm_global-scale": [
            f"model.layers.{i}.input_layernorm.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-mlp_global-wo-kernel": [
            f"model.layers.{i}.mlp.down_proj.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-mlp_global-wi_1-kernel": [
            f"model.layers.{i}.mlp.up_proj.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-mlp_global-wi_0-kernel": [
            f"model.layers.{i}.mlp.gate_proj.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-post_self_attention_norm_global-scale": [
            f"model.layers.{i}.post_attention_layernorm.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-post_ffw_norm_global-scale": [
            f"model.layers.{i}.post_feedforward_layernorm.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-pre_ffw_norm_global-scale": [
            f"model.layers.{i}.pre_feedforward_layernorm.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-self_attention_global-key-kernel": [
            f"model.layers.{i}.self_attn.k_proj.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-self_attention_global-out-kernel": [
            f"model.layers.{i}.self_attn.o_proj.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-self_attention_global-query-kernel": [
            f"model.layers.{i}.self_attn.q_proj.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-self_attention_global-value-kernel": [
            f"model.layers.{i}.self_attn.v_proj.weight" for i in range(1, nlayers, 2)
        ],
        "params-decoder-layers-pre_self_attention_norm_local-scale": [
            f"model.layers.{i}.input_layernorm.weight" for i in range(0, nlayers, 2)
        ],
        "params-decoder-layers-mlp_local-wo-kernel": [
            f"model.layers.{i}.mlp.down_proj.weight" for i in range(0, nlayers, 2)
        ],
        "params-decoder-layers-mlp_local-wi_1-kernel": [
            f"model.layers.{i}.mlp.up_proj.weight" for i in range(0, nlayers, 2)
        ],
        "params-decoder-layers-mlp_local-wi_0-kernel": [
            f"model.layers.{i}.mlp.gate_proj.weight" for i in range(0, nlayers, 2)
        ],
        "params-decoder-layers-post_self_attention_norm_local-scale": [
            f"model.layers.{i}.post_attention_layernorm.weight" for i in range(0, nlayers, 2)
        ],
        "params-decoder-layers-post_ffw_norm_local-scale": [
            f"model.layers.{i}.post_feedforward_layernorm.weight" for i in range(0, nlayers, 2)
        ],
        "params-decoder-layers-pre_ffw_norm_local-scale": [
            f"model.layers.{i}.pre_feedforward_layernorm.weight" for i in range(0, nlayers, 2)
        ],
        "params-decoder-layers-self_attention_local-key-kernel": [
            f"model.layers.{i}.self_attn.k_proj.weight" for i in range(0, nlayers, 2)
        ],
        "params-decoder-layers-self_attention_local-out-kernel": [
            f"model.layers.{i}.self_attn.o_proj.weight" for i in range(0, nlayers, 2)
        ],
        "params-decoder-layers-self_attention_local-query-kernel": [
            f"model.layers.{i}.self_attn.q_proj.weight" for i in range(0, nlayers, 2)
        ],
        "params-decoder-layers-self_attention_local-value-kernel": [
            f"model.layers.{i}.self_attn.v_proj.weight" for i in range(0, nlayers, 2)
        ],
    }
  # Case 2: scan_layer=False
  else:
    for maxtext_layer_idx in range(0, nlayers // 2):
      local_layer_idx = maxtext_layer_idx * 2
      global_layer_idx = maxtext_layer_idx * 2 + 1
      # pylint: disable=line-too-long
      layer_mapping = {
          f"params-decoder-layers_{maxtext_layer_idx}-pre_self_attention_norm_global-scale": f"model.layers.{global_layer_idx}.input_layernorm.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_global-wo-kernel": f"model.layers.{global_layer_idx}.mlp.down_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_global-wi_1-kernel": f"model.layers.{global_layer_idx}.mlp.up_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_global-wi_0-kernel": f"model.layers.{global_layer_idx}.mlp.gate_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-post_self_attention_norm_global-scale": f"model.layers.{global_layer_idx}.post_attention_layernorm.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-post_ffw_norm_global-scale": f"model.layers.{global_layer_idx}.post_feedforward_layernorm.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-pre_ffw_norm_global-scale": f"model.layers.{global_layer_idx}.pre_feedforward_layernorm.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_global-key-kernel": f"model.layers.{global_layer_idx}.self_attn.k_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_global-out-kernel": f"model.layers.{global_layer_idx}.self_attn.o_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_global-query-kernel": f"model.layers.{global_layer_idx}.self_attn.q_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_global-value-kernel": f"model.layers.{global_layer_idx}.self_attn.v_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-pre_self_attention_norm_local-scale": f"model.layers.{local_layer_idx}.input_layernorm.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_local-wo-kernel": f"model.layers.{local_layer_idx}.mlp.down_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_local-wi_1-kernel": f"model.layers.{local_layer_idx}.mlp.up_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_local-wi_0-kernel": f"model.layers.{local_layer_idx}.mlp.gate_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-post_self_attention_norm_local-scale": f"model.layers.{local_layer_idx}.post_attention_layernorm.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-post_ffw_norm_local-scale": f"model.layers.{local_layer_idx}.post_feedforward_layernorm.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-pre_ffw_norm_local-scale": f"model.layers.{local_layer_idx}.pre_feedforward_layernorm.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_local-key-kernel": f"model.layers.{local_layer_idx}.self_attn.k_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_local-out-kernel": f"model.layers.{local_layer_idx}.self_attn.o_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_local-query-kernel": f"model.layers.{local_layer_idx}.self_attn.q_proj.weight",
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_local-value-kernel": f"model.layers.{local_layer_idx}.self_attn.v_proj.weight",
      }
      mapping = {**mapping, **layer_mapping}
  return mapping


def GEMMA2_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
  """Creates parameter transformation functions for Gemma2 conversion.

  This function generates a mapping of transformation functions that handle the
  necessary conversions between MaxText and HuggingFace parameter formats for
  Gemma2, including operations like padding, reshaping, and scaling.

  Args:
      config (dict): Model configuration dictionary that must contain:
          - num_hidden_layers (int): Number of layers in the model.
          - head_dim (int): Dimension of attention heads.
          - hidden_size (int): Model's hidden dimension size.
      scan_layers (bool, optional): Controls the output format for layer
        parameters. True for batched, False for individual. Defaults to False.
      saving_to_hf (bool, optional): Determines the direction of transformation.
        True for MaxText to HuggingFace, False for the reverse. Defaults to
        False.

  Returns:
      dict: A mapping from MaxText parameter names to transformation functions.
        The value can be a single function or a list of functions to be
        applied sequentially.
  """
  nlayers = config["num_hidden_layers"]

  def pad_hf_embedding_layer(input_tensor, target_shape):
    """Pads/unpads and scales the embedding layer.

    Note:
        HF embedding weights shape =  [256000, d_model]
        MaxText embedding weights shape = [256128, d_model]
        MaxText pads Gemma2 embedding to 256128 for better performance.
    """
    # TODO(wenxindongwork), Perhaps, this dtype should be the activation dtype
    normalizer = np.dtype("float32").type(config["hidden_size"] ** 0.5)

    if saving_to_hf:
      target_tensor = input_tensor[: target_shape[0], : target_shape[1]]
      target_tensor = target_tensor / normalizer
      target_tensor = target_tensor.astype(input_tensor.dtype)
      return target_tensor
    else:
      target_tensor = np.zeros(target_shape, dtype=input_tensor.dtype)
      target_tensor[: input_tensor.shape[0], : input_tensor.shape[1]] = input_tensor
      target_tensor = target_tensor * normalizer
      target_tensor = target_tensor.astype(input_tensor.dtype)
      return target_tensor

  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
    else:
      return input_tensor.T.reshape(target_shape)

  def scale_rmsnorm_layer(input_tensor, target_shape):
    if saving_to_hf:
      return (input_tensor - 1.0).reshape(target_shape)
    else:
      return (input_tensor + 1.0).reshape(target_shape)

  def scale_query_layer(input_tensor, target_shape):
    if saving_to_hf:
      depth_scale = np.dtype("float32").type(np.sqrt(config["head_dim"]))
      return (input_tensor * depth_scale).astype(input_tensor.dtype)
    else:
      depth_scale = np.dtype("float32").type(1 / np.sqrt(config["head_dim"]))
      return (input_tensor * depth_scale).astype(input_tensor.dtype)

  # hook order does not affect result
  query_hook_chain = [reshape_kernel, scale_query_layer]

  mapping = {
      "params-token_embedder-embedding": pad_hf_embedding_layer,
      "params-decoder-decoder_norm-scale": scale_rmsnorm_layer,
  }
  if scan_layers:
    mapping = {
        **mapping,
        "params-decoder-layers-self_attention_global-query-kernel": query_hook_chain,
        "params-decoder-layers-self_attention_local-query-kernel": query_hook_chain,
        "params-decoder-layers-self_attention_global-key-kernel": reshape_kernel,
        "params-decoder-layers-self_attention_local-key-kernel": reshape_kernel,
        "params-decoder-layers-self_attention_global-value-kernel": reshape_kernel,
        "params-decoder-layers-self_attention_local-value-kernel": reshape_kernel,
        "params-decoder-layers-mlp_global-wo-kernel": reshape_kernel,
        "params-decoder-layers-mlp_global-wi_1-kernel": reshape_kernel,
        "params-decoder-layers-mlp_global-wi_0-kernel": reshape_kernel,
        "params-decoder-layers-self_attention_global-out-kernel": reshape_kernel,
        "params-decoder-layers-mlp_local-wo-kernel": reshape_kernel,
        "params-decoder-layers-mlp_local-wi_1-kernel": reshape_kernel,
        "params-decoder-layers-mlp_local-wi_0-kernel": reshape_kernel,
        "params-decoder-layers-self_attention_local-out-kernel": reshape_kernel,
        "params-decoder-layers-pre_self_attention_norm_global-scale": scale_rmsnorm_layer,
        "params-decoder-layers-post_self_attention_norm_global-scale": scale_rmsnorm_layer,
        "params-decoder-layers-post_ffw_norm_global-scale": scale_rmsnorm_layer,
        "params-decoder-layers-pre_ffw_norm_global-scale": scale_rmsnorm_layer,
        "params-decoder-layers-pre_self_attention_norm_local-scale": scale_rmsnorm_layer,
        "params-decoder-layers-post_self_attention_norm_local-scale": scale_rmsnorm_layer,
        "params-decoder-layers-post_ffw_norm_local-scale": scale_rmsnorm_layer,
        "params-decoder-layers-pre_ffw_norm_local-scale": scale_rmsnorm_layer,
    }
  else:
    for maxtext_layer_idx in range(nlayers // 2):
      mapping = {
          **mapping,
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_global-query-kernel": query_hook_chain,
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_local-query-kernel": query_hook_chain,
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_global-key-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_local-key-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_global-value-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_local-value-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_global-wo-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_global-wi_1-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_global-wi_0-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_global-out-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_local-wo-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_local-wi_1-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-mlp_local-wi_0-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-self_attention_local-out-kernel": reshape_kernel,
          f"params-decoder-layers_{maxtext_layer_idx}-pre_self_attention_norm_global-scale": scale_rmsnorm_layer,
          f"params-decoder-layers_{maxtext_layer_idx}-post_self_attention_norm_global-scale": scale_rmsnorm_layer,
          f"params-decoder-layers_{maxtext_layer_idx}-post_ffw_norm_global-scale": scale_rmsnorm_layer,
          f"params-decoder-layers_{maxtext_layer_idx}-pre_ffw_norm_global-scale": scale_rmsnorm_layer,
          f"params-decoder-layers_{maxtext_layer_idx}-pre_self_attention_norm_local-scale": scale_rmsnorm_layer,
          f"params-decoder-layers_{maxtext_layer_idx}-post_self_attention_norm_local-scale": scale_rmsnorm_layer,
          f"params-decoder-layers_{maxtext_layer_idx}-post_ffw_norm_local-scale": scale_rmsnorm_layer,
          f"params-decoder-layers_{maxtext_layer_idx}-pre_ffw_norm_local-scale": scale_rmsnorm_layer,
      }
  return mapping


def QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
  """Returns mapping from MaxText to HuggingFace Qwen weight paths.

  This function generates a dictionary that maps parameter names from a MaxText
  Qwen checkpoint to their corresponding names in the Hugging Face format.
  It handles both dense and Mixture-of-Experts (MoE) model variants.

  Args:
    config (dict): Model configuration dictionary, including
      'num_hidden_layers' and optionally 'num_experts'.
    scan_layers (bool, optional): Whether the MaxText model uses scanned
      layers. Defaults to False.

  Returns:
    dict: A mapping where keys are `atomic_mt_key` (single MaxText parameter names).
      Values are Hugging Face parameter names in one of four forms: unscanned (string),
      scanned (list of strings), unscanned with expert stacking (list of strings),
      or scanned with expert stacking (nested list of strings).
  """
  n_layers = config["num_hidden_layers"]
  num_experts = config.get("num_experts", 0)

  mapping = {
      "params-token_embedder-embedding": "model.embed_tokens.weight",
      "params-decoder-decoder_norm-scale": "model.norm.weight",
      "params-decoder-logits_dense-kernel": "lm_head.weight",
  }

  if scan_layers:
    # This block handles scanned layers for both dense and MoE models.
    mapping.update(
        {
            "params-decoder-layers-pre_self_attention_layer_norm-scale": [
                f"model.layers.{i}.input_layernorm.weight" for i in range(n_layers)
            ],
            "params-decoder-layers-self_attention-query-kernel": [
                f"model.layers.{i}.self_attn.q_proj.weight" for i in range(n_layers)
            ],
            "params-decoder-layers-self_attention-key-kernel": [
                f"model.layers.{i}.self_attn.k_proj.weight" for i in range(n_layers)
            ],
            "params-decoder-layers-self_attention-value-kernel": [
                f"model.layers.{i}.self_attn.v_proj.weight" for i in range(n_layers)
            ],
            "params-decoder-layers-self_attention-query-bias": [
                f"model.layers.{i}.self_attn.q_proj.bias" for i in range(n_layers)
            ],
            "params-decoder-layers-self_attention-key-bias": [
                f"model.layers.{i}.self_attn.k_proj.bias" for i in range(n_layers)
            ],
            "params-decoder-layers-self_attention-value-bias": [
                f"model.layers.{i}.self_attn.v_proj.bias" for i in range(n_layers)
            ],
            "params-decoder-layers-self_attention-out-kernel": [
                f"model.layers.{i}.self_attn.o_proj.weight" for i in range(n_layers)
            ],
            "params-decoder-layers-self_attention-query_norm-scale": [
                f"model.layers.{i}.self_attn.q_norm.weight" for i in range(n_layers)
            ],
            "params-decoder-layers-self_attention-key_norm-scale": [
                f"model.layers.{i}.self_attn.k_norm.weight" for i in range(n_layers)
            ],
            "params-decoder-layers-post_self_attention_layer_norm-scale": [
                f"model.layers.{i}.post_attention_layernorm.weight" for i in range(n_layers)
            ],
        }
    )
    if num_experts > 1:
      # For scanned MoE, we create a nested list: [[e0_l0, e0_l1..], [e1_l0, e1_l1..]..]
      # This follows the (experts, layers, ...) tensor layout.
      mapping.update(
          {
              "params-decoder-layers-moe_block-gate-kernel": [
                  f"model.layers.{i}.mlp.gate.weight" for i in range(n_layers)
              ],
              "params-decoder-layers-moe_block-wi_0": [
                  [f"model.layers.{l}.mlp.experts.{e}.gate_proj.weight" for l in range(n_layers)]
                  for e in range(num_experts)
              ],
              "params-decoder-layers-moe_block-wi_1": [
                  [f"model.layers.{l}.mlp.experts.{e}.up_proj.weight" for l in range(n_layers)]
                  for e in range(num_experts)
              ],
              "params-decoder-layers-moe_block-wo": [
                  [f"model.layers.{l}.mlp.experts.{e}.down_proj.weight" for l in range(n_layers)]
                  for e in range(num_experts)
              ],
          }
      )
    else:  # Dense MLP
      mapping.update(
          {
              "params-decoder-layers-mlp-wi_0-kernel": [
                  f"model.layers.{i}.mlp.gate_proj.weight" for i in range(n_layers)
              ],
              "params-decoder-layers-mlp-wi_1-kernel": [f"model.layers.{i}.mlp.up_proj.weight" for i in range(n_layers)],
              "params-decoder-layers-mlp-wo-kernel": [f"model.layers.{i}.mlp.down_proj.weight" for i in range(n_layers)],
          }
      )
  else:  # unscanned layers
    for i in range(n_layers):
      # Common Attention and Norms
      # pylint: disable=line-too-long
      mapping.update(
          {
              f"params-decoder-layers_{i}-pre_self_attention_layer_norm-scale": f"model.layers.{i}.input_layernorm.weight",
              f"params-decoder-layers_{i}-self_attention-query-kernel": f"model.layers.{i}.self_attn.q_proj.weight",
              f"params-decoder-layers_{i}-self_attention-key-kernel": f"model.layers.{i}.self_attn.k_proj.weight",
              f"params-decoder-layers_{i}-self_attention-value-kernel": f"model.layers.{i}.self_attn.v_proj.weight",
              f"params-decoder-layers_{i}-self_attention-out-kernel": f"model.layers.{i}.self_attn.o_proj.weight",
              f"params-decoder-layers_{i}-self_attention-query-bias": f"model.layers.{i}.self_attn.q_proj.bias",
              f"params-decoder-layers_{i}-self_attention-key-bias": f"model.layers.{i}.self_attn.k_proj.bias",
              f"params-decoder-layers_{i}-self_attention-value-bias": f"model.layers.{i}.self_attn.v_proj.bias",
              f"params-decoder-layers_{i}-self_attention-query_norm-scale": f"model.layers.{i}.self_attn.q_norm.weight",
              f"params-decoder-layers_{i}-self_attention-key_norm-scale": f"model.layers.{i}.self_attn.k_norm.weight",
              f"params-decoder-layers_{i}-post_self_attention_layer_norm-scale": f"model.layers.{i}.post_attention_layernorm.weight",
              f"params-decoder-layers_{i}-post_self_attention_layer_norm-scale": f"model.layers.{i}.post_attention_layernorm.weight",
          }
      )
      if num_experts > 1:
        # For each unscanned MoE layer, map the MaxText parameter to a 1D list of all expert weights for that layer.
        mapping.update(
            {
                f"params-decoder-layers_{i}-moe_block-gate-kernel": f"model.layers.{i}.mlp.gate.weight",
                f"params-decoder-layers_{i}-moe_block-wi_0": [
                    f"model.layers.{i}.mlp.experts.{j}.gate_proj.weight" for j in range(num_experts)
                ],
                f"params-decoder-layers_{i}-moe_block-wi_1": [
                    f"model.layers.{i}.mlp.experts.{j}.up_proj.weight" for j in range(num_experts)
                ],
                f"params-decoder-layers_{i}-moe_block-wo": [
                    f"model.layers.{i}.mlp.experts.{j}.down_proj.weight" for j in range(num_experts)
                ],
            }
        )
      else:  # Dense MLP
        mapping.update(
            {
                f"params-decoder-layers_{i}-mlp-wi_0-kernel": f"model.layers.{i}.mlp.gate_proj.weight",
                f"params-decoder-layers_{i}-mlp-wi_1-kernel": f"model.layers.{i}.mlp.up_proj.weight",
                f"params-decoder-layers_{i}-mlp-wo-kernel": f"model.layers.{i}.mlp.down_proj.weight",
            }
        )
  return mapping


def QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
  """Creates parameter transformation functions for Qwen.

  This function provides a dictionary of transformation functions (hooks) for
  converting Qwen model parameters between MaxText and Hugging Face formats.
  It handles embedding padding and kernel reshaping.

  Args:
    config (dict): Model configuration dictionary, including
      'num_hidden_layers' and optionally 'num_experts'.
    scan_layers (bool, optional): Whether the model uses scanned layers.
      Defaults to False.
    saving_to_hf (bool, optional): The direction of conversion. True for
      MaxText to Hugging Face, False for the reverse. Defaults to False.

  Returns:
    dict: A dictionary mapping MaxText parameter names to their corresponding
      transformation functions.
  """
  n_layers = config["num_hidden_layers"]
  num_experts = config.get("num_experts", 0)

  def pad_embedding_layer(input_tensor, target_shape):
    """Pads or truncates embedding layer to match target vocab size."""
    source_vocab_size = input_tensor.shape[0]
    target_vocab_size = target_shape[0]

    if source_vocab_size == target_vocab_size:
      return input_tensor

    if saving_to_hf:  # MaxText to HF, truncate
      return input_tensor[:target_vocab_size, :]
    else:  # HF to MaxText, pad
      padded_tensor = np.zeros(target_shape, dtype=input_tensor.dtype)
      padded_tensor[:source_vocab_size, :] = input_tensor
      return padded_tensor

  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
    else:
      return input_tensor.T.reshape(target_shape)

  def reshape_bias(input_tensor, target_shape=None):
    """Reshapes biases between MaxText 2D (heads, dim) and HF 1D (hidden)."""
    # saving_to_hf: MaxText [heads, head_dim] -> HF [hidden_dim] (flatten)
    # loading_to_maxtext: HF [hidden_dim] -> MaxText [heads, head_dim]
    return input_tensor.reshape(target_shape)

  mapping = {
      "params-token_embedder-embedding": pad_embedding_layer,
      "params-decoder-logits_dense-kernel": reshape_kernel,
  }

  kernel_hooks = [
      "self_attention-query-kernel",
      "self_attention-key-kernel",
      "self_attention-value-kernel",
      "self_attention-out-kernel",
      "mlp-wi_0-kernel",
      "mlp-wi_1-kernel",
      "mlp-wo-kernel",
  ]
  bias_hooks = [
      "self_attention-query-bias",
      "self_attention-key-bias",
      "self_attention-value-bias",
  ]
  moe_kernel_hooks = [
      "moe_block-gate-kernel",
      "moe_block-wi_0-kernel",
      "moe_block-wi_1-kernel",
      "moe_block-wo-kernel",
      "moe_block-wi_0",
      "moe_block-wi_1",
      "moe_block-wo",
  ]

  if scan_layers:
    for key in kernel_hooks:
      mapping[f"params-decoder-layers-{key}"] = reshape_kernel
    for key in bias_hooks:
      mapping[f"params-decoder-layers-{key}"] = reshape_bias
    if num_experts > 1:
      for key in moe_kernel_hooks:
        mapping[f"params-decoder-layers-{key}"] = reshape_kernel
  else:
    for i in range(n_layers):
      for key in kernel_hooks:
        mapping[f"params-decoder-layers_{i}-{key}"] = reshape_kernel
      for key in bias_hooks:
        mapping[f"params-decoder-layers_{i}-{key}"] = reshape_bias
      if num_experts > 1:
        for key in moe_kernel_hooks:
          mapping[f"params-decoder-layers_{i}-{key}"] = reshape_kernel
  return mapping


def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
  """
  Returns mapping from MaxText to HuggingFace Qwen3-Next weight paths.
  All MaxText keys start with 'params-' and use '-' separators for scanned layers.
  """
  num_main_layers = config["num_hidden_layers"]
  num_experts = config["num_experts"]
  layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval

  # 1. Non-layer specific weight mappings
  mapping = {
      "params-token_embedder-embedding": "model.embed_tokens.weight",
      "params-decoder-decoder_norm-scale": "model.norm.weight",
      "params-decoder-logits_dense-kernel": "lm_head.weight",
  }

  if scan_layers:
    # 2. Scan over block cycles
    for block_idx in range(layer_cycle_interval):
      hf_indices = list(range(block_idx, num_main_layers, layer_cycle_interval))
      prefix = f"params-decoder-layers-layer_{block_idx}"

      # Layer norms
      mapping[f"{prefix}-input_layernorm-scale"] = [f"model.layers.{i}.input_layernorm.weight" for i in hf_indices]
      mapping[f"{prefix}-post_attention_layernorm-scale"] = [
          f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices
      ]

      # Handle Interleaved Attention (Linear vs Full)
      is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0

      if is_full_attention_layer:
        mapping.update(
            {
                f"{prefix}-attention-attention-query-kernel": [
                    f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices
                ],
                f"{prefix}-attention-attention-key-kernel": [
                    f"model.layers.{i}.self_attn.k_proj.weight" for i in hf_indices
                ],
                f"{prefix}-attention-attention-value-kernel": [
                    f"model.layers.{i}.self_attn.v_proj.weight" for i in hf_indices
                ],
                f"{prefix}-attention-attention-out-kernel": [
                    f"model.layers.{i}.self_attn.o_proj.weight" for i in hf_indices
                ],
                f"{prefix}-attention-attention-query_norm-scale": [
                    f"model.layers.{i}.self_attn.q_norm.weight" for i in hf_indices
                ],
                f"{prefix}-attention-attention-key_norm-scale": [
                    f"model.layers.{i}.self_attn.k_norm.weight" for i in hf_indices
                ],
            }
        )
      else:
        # Linear/Hybrid Attention Block
        mapping.update(
            {
                f"{prefix}-attention-in_proj_qkvz-kernel": [
                    f"model.layers.{i}.linear_attn.in_proj_qkvz.weight" for i in hf_indices
                ],
                f"{prefix}-attention-in_proj_ba-kernel": [
                    f"model.layers.{i}.linear_attn.in_proj_ba.weight" for i in hf_indices
                ],
                f"{prefix}-attention-conv1d-kernel": [f"model.layers.{i}.linear_attn.conv1d.weight" for i in hf_indices],
                f"{prefix}-attention-A_log": [f"model.layers.{i}.linear_attn.A_log" for i in hf_indices],
                f"{prefix}-attention-dt_bias": [f"model.layers.{i}.linear_attn.dt_bias" for i in hf_indices],
                f"{prefix}-attention-norm-rms_norm-scale": [
                    f"model.layers.{i}.linear_attn.norm.weight" for i in hf_indices
                ],
                f"{prefix}-attention-out_proj-kernel": [
                    f"model.layers.{i}.linear_attn.out_proj.weight" for i in hf_indices
                ],
            }
        )

      # 3. Handle MLP: Gates and Shared Experts
      mapping.update(
          {
              f"{prefix}-mlp-routed_experts-gate-kernel": [f"model.layers.{i}.mlp.gate.weight" for i in hf_indices],
              f"{prefix}-mlp-shared_expert-wi_0-kernel": [
                  f"model.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices
              ],
              f"{prefix}-mlp-shared_expert-wi_1-kernel": [
                  f"model.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices
              ],
              f"{prefix}-mlp-shared_expert-wo-kernel": [
                  f"model.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices
              ],
              f"{prefix}-mlp-shared_expert_gate-kernel": [
                  f"model.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices
              ],
          }
      )

      # 4. Handle MoE Routed Experts
      mapping.update(
          {
              f"{prefix}-mlp-routed_experts-wi_0": [
                  [f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices] for e in range(num_experts)
              ],
              f"{prefix}-mlp-routed_experts-wi_1": [
                  [f"model.layers.{i}.mlp.experts.{e}.up_proj.weight" for i in hf_indices] for e in range(num_experts)
              ],
              f"{prefix}-mlp-routed_experts-wo": [
                  [f"model.layers.{i}.mlp.experts.{e}.down_proj.weight" for i in hf_indices] for e in range(num_experts)
              ],
          }
      )
  else:
    # Unscanned layer mapping
    for i in range(num_main_layers):
      prefix = f"params-decoder-layers_{i}"

      # Layer Norms
      mapping[f"{prefix}-input_layernorm-scale"] = f"model.layers.{i}.input_layernorm.weight"
      mapping[f"{prefix}-post_attention_layernorm-scale"] = f"model.layers.{i}.post_attention_layernorm.weight"

      # Determine layer type based on cycle interval
      # Assuming block logic: layer i corresponds to block_idx = i % interval
      block_idx = i % layer_cycle_interval
      is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0

      if is_full_attention_layer:
        mapping.update(
            {
                f"{prefix}-attention-attention-query-kernel": f"model.layers.{i}.self_attn.q_proj.weight",
                f"{prefix}-attention-attention-key-kernel": f"model.layers.{i}.self_attn.k_proj.weight",
                f"{prefix}-attention-attention-value-kernel": f"model.layers.{i}.self_attn.v_proj.weight",
                f"{prefix}-attention-attention-out-kernel": f"model.layers.{i}.self_attn.o_proj.weight",
                f"{prefix}-attention-attention-query_norm-scale": f"model.layers.{i}.self_attn.q_norm.weight",
                f"{prefix}-attention-attention-key_norm-scale": f"model.layers.{i}.self_attn.k_norm.weight",
            }
        )
      else:
        # Linear/Hybrid Attention Block
        mapping.update(
            {
                f"{prefix}-attention-in_proj_qkvz-kernel": f"model.layers.{i}.linear_attn.in_proj_qkvz.weight",
                f"{prefix}-attention-in_proj_ba-kernel": f"model.layers.{i}.linear_attn.in_proj_ba.weight",
                f"{prefix}-attention-conv1d-kernel": f"model.layers.{i}.linear_attn.conv1d.weight",
                f"{prefix}-attention-A_log": f"model.layers.{i}.linear_attn.A_log",
                f"{prefix}-attention-dt_bias": f"model.layers.{i}.linear_attn.dt_bias",
                f"{prefix}-attention-norm-rms_norm-scale": f"model.layers.{i}.linear_attn.norm.weight",
                f"{prefix}-attention-out_proj-kernel": f"model.layers.{i}.linear_attn.out_proj.weight",
            }
        )

      # MLP: Gates and Shared Experts
      mapping.update(
          {
              f"{prefix}-mlp-routed_experts-gate-kernel": f"model.layers.{i}.mlp.gate.weight",
              f"{prefix}-mlp-shared_expert-wi_0-kernel": f"model.layers.{i}.mlp.shared_expert.gate_proj.weight",
              f"{prefix}-mlp-shared_expert-wi_1-kernel": f"model.layers.{i}.mlp.shared_expert.up_proj.weight",
              f"{prefix}-mlp-shared_expert-wo-kernel": f"model.layers.{i}.mlp.shared_expert.down_proj.weight",
              f"{prefix}-mlp-shared_expert_gate-kernel": f"model.layers.{i}.mlp.shared_expert_gate.weight",
          }
      )

      # MoE Routed Experts (List of expert weights for this specific layer)
      mapping.update(
          {
              f"{prefix}-mlp-routed_experts-wi_0": [
                  f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for e in range(num_experts)
              ],
              f"{prefix}-mlp-routed_experts-wi_1": [
                  f"model.layers.{i}.mlp.experts.{e}.up_proj.weight" for e in range(num_experts)
              ],
              f"{prefix}-mlp-routed_experts-wo": [
                  f"model.layers.{i}.mlp.experts.{e}.down_proj.weight" for e in range(num_experts)
              ],
          }
      )
  return mapping


def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
  """
  Transformation hooks for parameters using hyphenated 'params-' MaxText keys.
  """

  def transpose(input_tensor, target_shape=None):
    return input_tensor.T

  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
    else:
      return input_tensor.T.reshape(target_shape)

  def permute_conv(input_tensor, target_shape=None):
    # MT: [K, 1, C] <-> HF: [C, 1, K]
    return input_tensor.transpose(2, 1, 0)

  # Initialize Hooks
  hooks = {
      "params-decoder-logits_dense-kernel": transpose,
  }

  layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval
  num_main_layers = config["num_hidden_layers"]
  loop_indices = range(layer_cycle_interval) if scan_layers else range(num_main_layers)

  for i in loop_indices:
    if scan_layers:
      prefix = f"params-decoder-layers-layer_{i}"
      block_idx = i
    else:
      prefix = f"params-decoder-layers_{i}"
      block_idx = i % layer_cycle_interval
    is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0

    if is_full_attention_layer:
      for key in ["query", "key", "value", "out"]:
        hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel
    else:
      hooks[f"{prefix}-attention-in_proj_qkvz-kernel"] = transpose
      hooks[f"{prefix}-attention-in_proj_ba-kernel"] = transpose
      hooks[f"{prefix}-attention-out_proj-kernel"] = transpose
      hooks[f"{prefix}-attention-conv1d-kernel"] = permute_conv

    mlp_prefix = f"{prefix}-mlp"
    hooks[f"{mlp_prefix}-routed_experts-gate-kernel"] = transpose
    hooks[f"{mlp_prefix}-shared_expert-wi_0-kernel"] = transpose
    hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel"] = transpose
    hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose
    hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose

    hooks[f"{mlp_prefix}-routed_experts-wi_0"] = transpose
    hooks[f"{mlp_prefix}-routed_experts-wi_1"] = transpose
    hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose

  return hooks


def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
  """Generates a parameter mapping from MaxText to HuggingFace Deepseek weight paths.

  Returns:
    dict: A mapping where keys are `atomic_mt_key` (single MaxText parameter names).
      Values are Hugging Face parameter names in one of four forms: unscanned (string),
      scanned (list of strings), unscanned with expert stacking (list of strings),
      or scanned with expert stacking (nested list of strings).
  """
  # TODO(shuningjin): add unscan support, b/457820735
  if not scan_layers:
    raise NotImplementedError("This conversion only supports scanned MaxText models.")

  # Extract hf configuration parameters, without mtp
  num_main_layers = config["num_hidden_layers"]
  first_num_dense_layers = config["first_k_dense_replace"]
  num_experts = config.get("n_routed_experts", 0)

  # Mapping for non-layer-specific weights
  mapping = {
      "params-token_embedder-embedding": "model.embed_tokens.weight",
      "params-decoder-decoder_norm-scale": "model.norm.weight",
      "params-decoder-logits_dense-kernel": "lm_head.weight",
  }
  # Attention keys are shared by both dense and MoE
  attention_keys = {
      "pre_self_attention_layer_norm-scale": "input_layernorm.weight",
      "post_self_attention_layer_norm-scale": "post_attention_layernorm.weight",
      "self_attention-wkv_a-kernel": "self_attn.kv_a_proj_with_mqa.weight",
      "self_attention-kv_norm-scale": "self_attn.kv_a_layernorm.weight",
      "self_attention-wkv_b-kernel": "self_attn.kv_b_proj.weight",
      "self_attention-out-kernel": "self_attn.o_proj.weight",
      # v3
      "self_attention-wq_a-kernel": "self_attn.q_a_proj.weight",
      "self_attention-q_norm-scale": "self_attn.q_a_layernorm.weight",
      "self_attention-wq_b-kernel": "self_attn.q_b_proj.weight",
      # v2
      "self_attention-query-kernel": "self_attn.q_proj.weight",
  }
  # Dense Layers
  dense_layer_keys = attention_keys | {
      "mlp-wi_0-kernel": "mlp.gate_proj.weight",
      "mlp-wi_1-kernel": "mlp.up_proj.weight",
      "mlp-wo-kernel": "mlp.down_proj.weight",
  }
  for maxtext_key, hf_key in dense_layer_keys.items():
    mapping[f"params-decoder-dense_layers-{maxtext_key}"] = [
        f"model.layers.{i}.{hf_key}" for i in range(first_num_dense_layers)
    ]

  # MoE Layers
  moe_layer_keys = attention_keys | {
      "DeepSeekMoeBlock_0-shared_experts-wi_0-kernel": "mlp.shared_experts.gate_proj.weight",
      "DeepSeekMoeBlock_0-shared_experts-wi_1-kernel": "mlp.shared_experts.up_proj.weight",
      "DeepSeekMoeBlock_0-shared_experts-wo-kernel": "mlp.shared_experts.down_proj.weight",
      "DeepSeekMoeBlock_0-MoeBlock_0-gate-kernel": "mlp.gate.weight",
      # v3
      "DeepSeekMoeBlock_0-MoeBlock_0-gate-bias": "mlp.gate.e_score_correction_bias",
  }
  for maxtext_key, hf_key in moe_layer_keys.items():
    mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [
        f"model.layers.{i}.{hf_key}" for i in range(first_num_dense_layers, num_main_layers)
    ]

  # MoE Experts (nested list mapping: [[e0_l0, e0_l1..], [e1_l0, e1_l1..]..])
  moe_expert_keys = {
      "DeepSeekMoeBlock_0-MoeBlock_0-wi_0": "gate_proj.weight",
      "DeepSeekMoeBlock_0-MoeBlock_0-wi_1": "up_proj.weight",
      "DeepSeekMoeBlock_0-MoeBlock_0-wo": "down_proj.weight",
  }
  for maxtext_key, hf_key in moe_expert_keys.items():
    mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [
        [f"model.layers.{l}.mlp.experts.{e}.{hf_key}" for l in range(first_num_dense_layers, num_main_layers)]
        for e in range(num_experts)
    ]
  return mapping


def DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
  """Creates parameter transformation functions for Deepseek."""
  # TODO(shuningjin): support hf->orbax(scan), b/457820372
  if not saving_to_hf:
    raise NotImplementedError("This conversion only supports saving_to_hf")
  # TODO(shuningjin): add unscan support, b/457820735
  if not scan_layers:
    raise NotImplementedError("This conversion only supports scanned MaxText models.")

  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
    else:
      return input_tensor.T.reshape(target_shape)

  mapping = {
      "params-decoder-logits_dense-kernel": reshape_kernel,
  }
  # all keys that need the reshape hook
  params_need_reshape = {
      # Dense Layers
      "params-decoder-dense_layers-self_attention-query-kernel",
      "params-decoder-dense_layers-self_attention-wq_a-kernel",
      "params-decoder-dense_layers-self_attention-wq_b-kernel",
      "params-decoder-dense_layers-self_attention-wkv_a-kernel",
      "params-decoder-dense_layers-self_attention-wkv_b-kernel",
      "params-decoder-dense_layers-self_attention-out-kernel",
      "params-decoder-dense_layers-mlp-wi_0-kernel",
      "params-decoder-dense_layers-mlp-wi_1-kernel",
      "params-decoder-dense_layers-mlp-wo-kernel",
      # MoE Layers
      "params-decoder-moe_layers-self_attention-query-kernel",
      "params-decoder-moe_layers-self_attention-wq_a-kernel",
      "params-decoder-moe_layers-self_attention-wq_b-kernel",
      "params-decoder-moe_layers-self_attention-wkv_a-kernel",
      "params-decoder-moe_layers-self_attention-wkv_b-kernel",
      "params-decoder-moe_layers-self_attention-out-kernel",
      "params-decoder-moe_layers-DeepSeekMoeBlock_0-shared_experts-wi_0-kernel",
      "params-decoder-moe_layers-DeepSeekMoeBlock_0-shared_experts-wi_1-kernel",
      "params-decoder-moe_layers-DeepSeekMoeBlock_0-shared_experts-wo-kernel",
      "params-decoder-moe_layers-DeepSeekMoeBlock_0-MoeBlock_0-gate-kernel",
      "params-decoder-moe_layers-DeepSeekMoeBlock_0-MoeBlock_0-wi_0",
      "params-decoder-moe_layers-DeepSeekMoeBlock_0-MoeBlock_0-wi_1",
      "params-decoder-moe_layers-DeepSeekMoeBlock_0-MoeBlock_0-wo",
  }

  for key in params_need_reshape:
    mapping[key] = reshape_kernel
  return mapping


def DEEPSEEK_NNX_TO_VLLM_PARAM_HOOK_FN():
  """Creates parameter transformation functions for Deepseek."""
  return {}


def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
  """Generates mapping from MaxText gpt-oss to Hugging Face weight paths.

  Returns:
    dict: A mapping where keys are `atomic_mt_key` (single MaxText parameter) or
    `composite_mt_key` (a tuple of MaxText parameters). Values are Hugging Face parameter
    names either a single string (unscanned form) or a list of strings (scanned form).

  Notes:
  - Handles the inhomogeneous scan block structure, based on `inhomogeneous_layer_cycle_interval`
  - Handles `composite_mt_key`: multiple MaxText keys map to HF key(s)
    - (GptOssMlp-wi_0, GptOssMlp-wi_1): mlp.experts.gate_up_proj
    - (GptOssMlp-wi_0_bias, GptOssMlp-wi_1_bias): mlp.experts.gate_up_proj_bias
  """
  n_layers = config["num_hidden_layers"]  # hf config
  layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval

  # Base mapping for non-layer parameters (targeting standard HF keys)
  mapping = {
      "params-token_embedder-embedding": "model.embed_tokens.weight",
      "params-decoder-decoder_norm-scale": "model.norm.weight",
      "params-decoder-logits_dense-kernel": "lm_head.weight",
  }

  if scan_layers:
    # Scan over blocks
    for block_idx in range(layer_cycle_interval):
      # Identify all original HF layer indices that collapse into this block
      hf_indices = range(block_idx, n_layers, layer_cycle_interval)
      prefix = f"params-decoder-layers-layers_{block_idx}"
      block_mapping = {
          # Layer Norms
          f"{prefix}-pre_self_attention_layer_norm-scale": [
              f"model.layers.{i}.input_layernorm.weight" for i in hf_indices
          ],
          f"{prefix}-post_self_attention_layer_norm-scale": [
              f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices
          ],
          # GptOssAttention
          f"{prefix}-GptOssAttention-query-kernel": [f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices],
          f"{prefix}-GptOssAttention-query-bias": [f"model.layers.{i}.self_attn.q_proj.bias" for i in hf_indices],
          f"{prefix}-GptOssAttention-key-kernel": [f"model.layers.{i}.self_attn.k_proj.weight" for i in hf_indices],
          f"{prefix}-GptOssAttention-key-bias": [f"model.layers.{i}.self_attn.k_proj.bias" for i in hf_indices],
          f"{prefix}-GptOssAttention-value-kernel": [f"model.layers.{i}.self_attn.v_proj.weight" for i in hf_indices],
          f"{prefix}-GptOssAttention-value-bias": [f"model.layers.{i}.self_attn.v_proj.bias" for i in hf_indices],
          f"{prefix}-GptOssAttention-out-kernel": [f"model.layers.{i}.self_attn.o_proj.weight" for i in hf_indices],
          f"{prefix}-GptOssAttention-out-bias": [f"model.layers.{i}.self_attn.o_proj.bias" for i in hf_indices],
          f"{prefix}-GptOssAttention-sinks": [f"model.layers.{i}.self_attn.sinks" for i in hf_indices],
          # GptOssMlp
          # 1. Gate/Router
          f"{prefix}-GptOssMlp-gate-kernel": [f"model.layers.{i}.mlp.router.weight" for i in hf_indices],
          f"{prefix}-GptOssMlp-gate-bias": [f"model.layers.{i}.mlp.router.bias" for i in hf_indices],
          # 2. Experts (Down Projection)
          f"{prefix}-GptOssMlp-wo": [f"model.layers.{i}.mlp.experts.down_proj" for i in hf_indices],
          f"{prefix}-GptOssMlp-wo_bias": [f"model.layers.{i}.mlp.experts.down_proj_bias" for i in hf_indices],
          # 3. Experts (Gate/Up Fused Projection)
          # `composite_mt_key`: Multiple MaxText keys map to HF key(s).
          (f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1"): [
              f"model.layers.{i}.mlp.experts.gate_up_proj" for i in hf_indices
          ],
          (f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias"): [
              f"model.layers.{i}.mlp.experts.gate_up_proj_bias" for i in hf_indices
          ],
      }
      mapping.update(block_mapping)

  else:
    # Unscan
    for i in range(n_layers):
      prefix = f"params-decoder-layers_{i}"
      layer_mapping = {
          # Layer Norms
          f"{prefix}-pre_self_attention_layer_norm-scale": f"model.layers.{i}.input_layernorm.weight",
          f"{prefix}-post_self_attention_layer_norm-scale": f"model.layers.{i}.post_attention_layernorm.weight",
          # GptOssAttention
          f"{prefix}-GptOssAttention-query-kernel": f"model.layers.{i}.self_attn.q_proj.weight",
          f"{prefix}-GptOssAttention-query-bias": f"model.layers.{i}.self_attn.q_proj.bias",
          f"{prefix}-GptOssAttention-key-kernel": f"model.layers.{i}.self_attn.k_proj.weight",
          f"{prefix}-GptOssAttention-key-bias": f"model.layers.{i}.self_attn.k_proj.bias",
          f"{prefix}-GptOssAttention-value-kernel": f"model.layers.{i}.self_attn.v_proj.weight",
          f"{prefix}-GptOssAttention-value-bias": f"model.layers.{i}.self_attn.v_proj.bias",
          f"{prefix}-GptOssAttention-out-kernel": f"model.layers.{i}.self_attn.o_proj.weight",
          f"{prefix}-GptOssAttention-out-bias": f"model.layers.{i}.self_attn.o_proj.bias",
          f"{prefix}-GptOssAttention-sinks": f"model.layers.{i}.self_attn.sinks",
          # GptOssMlp
          # 1. Gate/Router
          f"{prefix}-GptOssMlp-gate-kernel": f"model.layers.{i}.mlp.router.weight",
          f"{prefix}-GptOssMlp-gate-bias": f"model.layers.{i}.mlp.router.bias",
          # 2. Experts (Down Projection)
          f"{prefix}-GptOssMlp-wo": f"model.layers.{i}.mlp.experts.down_proj",
          f"{prefix}-GptOssMlp-wo_bias": f"model.layers.{i}.mlp.experts.down_proj_bias",
          # 3. Experts (Gate/Up Fused Projection)
          # `composite_mt_key`: Multiple MaxText keys map to HF key(s).
          (f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1"): f"model.layers.{i}.mlp.experts.gate_up_proj",
          (
              f"{prefix}-GptOssMlp-wi_0_bias",
              f"{prefix}-GptOssMlp-wi_1_bias",
          ): f"model.layers.{i}.mlp.experts.gate_up_proj_bias",
      }
      mapping.update(layer_mapping)

  return mapping


def GPT_OSS_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
  """Transformation hooks for gpt-oss parameters.

  Notes:
  - Handles the inhomogeneous scan block structure (inhomogeneous_layer_cycle_interval)
  - Handles `composite_mt_key` where multiple MaxText keys map to HF key(s)
    - (GptOssMlp-wi_0, GptOssMlp-wi_1): mlp.experts.gate_up_proj
    - (GptOssMlp-wi_0_bias, GptOssMlp-wi_1_bias): mlp.experts.gate_up_proj_bias
    - The composite keys are transformed via `interleave` function
  """

  def transpose(input_tensor, target_shape=None):
    return input_tensor.T

  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
    else:
      return input_tensor.T.reshape(target_shape)

  def reshape_bias(input_tensor, target_shape=None):
    """Reshapes biases between MaxText 2D (heads, dim) and HF 1D (hidden)."""
    if saving_to_hf:
      # MaxText [heads, head_dim] -> HF [hidden_dim] (flatten)
      return input_tensor.reshape(target_shape)
    else:
      # HF [hidden_dim] -> MaxText [heads, head_dim]
      return input_tensor.reshape(target_shape)

  def interleave(input_tensor, target_shape=None):
    """
    Handles `composite_mt_key`: maxtext (wi_0, wi_1) <-> hf (wi_0_1)
    - if saving_to_hf: (wi_0, wi_1) -> wi_0_1
      - input_tensor is a list of two tensors, tensor ORDER must be same as key order
      - return a single tensor
    - otherwise: wi_0_1 -> (wi_0, wi_1)
      - input_tensor is a single tensor
      - return two tensors stack at LAST index -1, tensor ORDER must be same as key order
    """
    if saving_to_hf:
      wi_0, wi_1 = input_tensor
      wi_0_1 = np.empty(target_shape, dtype=wi_0.dtype)
      wi_0_1[..., ::2] = wi_0
      wi_0_1[..., 1::2] = wi_1
      return wi_0_1
    else:
      wi_0_1 = input_tensor
      wi_0 = wi_0_1[..., ::2]
      wi_1 = wi_0_1[..., 1::2]
      return np.stack([wi_0, wi_1], axis=-1)

  n_layers = config["num_hidden_layers"]  # hf config
  layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval

  hooks = {"params-decoder-logits_dense-kernel": transpose}

  indices = range(layer_cycle_interval) if scan_layers else range(n_layers)
  for idx in indices:
    prefix = f"params-decoder-layers-layers_{idx}" if scan_layers else f"params-decoder-layers_{idx}"
    # Attention Kernels & Biases
    for key in ["query", "key", "value"]:
      hooks[f"{prefix}-GptOssAttention-{key}-kernel"] = reshape_kernel
      hooks[f"{prefix}-GptOssAttention-{key}-bias"] = reshape_bias
    hooks[f"{prefix}-GptOssAttention-out-kernel"] = reshape_kernel
    # MLP Kernels & Biases
    hooks[f"{prefix}-GptOssMlp-gate-kernel"] = transpose
    # `composite_mt_key`: A hook for combining multiple MaxText params.
    hooks[(f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1")] = interleave
    hooks[(f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias")] = interleave

  return hooks


def QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
  """Returns mapping from MaxText to HuggingFace Qwen3-Omni weight paths.

  This function combines mappings from different modalities (text, vision, audio, etc.)
  into a unified parameter mapping for the multi-modal Qwen3-Omni model.

  Args:
    config (dict): Model configuration dictionary containing modality-specific configs.
    scan_layers (bool, optional): Whether the model uses scanned layers. Defaults to False.

  Returns:
    dict: Combined mapping from all modalities.
  """
  # Collect all modality mappings
  mapping = {}

  # Text mapping with "thinker." prefix, reusing QWEN3-MOE mapping function
  num_experts_text = config["thinker_config"]["text_config"].get("num_experts", 0)
  n_layers_text = config["thinker_config"]["text_config"]["num_hidden_layers"]
  text_mapping = QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(
      config={"num_hidden_layers": n_layers_text, "num_experts": num_experts_text},
      maxtext_config=maxtext_config,
      scan_layers=scan_layers,
  )

  # Add "thinker." prefix to text mapping values
  def add_prefix_recursive(value):
    """Recursively add 'thinker.' prefix to strings, handling nested lists."""
    if isinstance(value, list):
      return [add_prefix_recursive(v) for v in value]
    else:
      return f"thinker.{value}"

  for key, value in text_mapping.items():
    text_mapping[key] = add_prefix_recursive(value)
  mapping.update(text_mapping)

  # Vision mapping
  vision_config = config["thinker_config"]["vision_config"]
  n_vision_layers = vision_config["depth"]

  # Vision patch embedding
  mapping["params-vision_encoder-Qwen3OmniMoeVisionEncoder_0-patch_embed-proj-kernel"] = (
      "thinker.visual.patch_embed.proj.weight"
  )
  mapping["params-vision_encoder-Qwen3OmniMoeVisionEncoder_0-patch_embed-proj-bias"] = (
      "thinker.visual.patch_embed.proj.bias"
  )

  # Vision positional embedding
  mapping["params-vision_encoder-Qwen3OmniMoeVisionEncoder_0-pos_embed_interpolate-pos_embed"] = (
      "thinker.visual.pos_embed.weight"
  )

  # Vision blocks (27 layers)
  for i in range(n_vision_layers):
    prefix = f"params-vision_encoder-Qwen3OmniMoeVisionEncoder_0-blocks_{i}"
    hf_prefix = f"thinker.visual.blocks.{i}"

    # Layer norms
    mapping[f"{prefix}-ln1-scale"] = f"{hf_prefix}.norm1.weight"
    mapping[f"{prefix}-ln1-bias"] = f"{hf_prefix}.norm1.bias"
    mapping[f"{prefix}-ln2-scale"] = f"{hf_prefix}.norm2.weight"
    mapping[f"{prefix}-ln2-bias"] = f"{hf_prefix}.norm2.bias"

    # Attention (HF has fused QKV, MaxText has separate Q/K/V)
    # We'll handle the split/fusion in the hook functions
    mapping[f"{prefix}-attn-attn-query-kernel"] = f"{hf_prefix}.attn.qkv.weight"
    mapping[f"{prefix}-attn-attn-query-bias"] = f"{hf_prefix}.attn.qkv.bias"
    mapping[f"{prefix}-attn-attn-key-kernel"] = f"{hf_prefix}.attn.qkv.weight"
    mapping[f"{prefix}-attn-attn-key-bias"] = f"{hf_prefix}.attn.qkv.bias"
    mapping[f"{prefix}-attn-attn-value-kernel"] = f"{hf_prefix}.attn.qkv.weight"
    mapping[f"{prefix}-attn-attn-value-bias"] = f"{hf_prefix}.attn.qkv.bias"
    mapping[f"{prefix}-attn-attn-out-kernel"] = f"{hf_prefix}.attn.proj.weight"
    mapping[f"{prefix}-attn-attn-out-bias"] = f"{hf_prefix}.attn.proj.bias"

    # MLP
    mapping[f"{prefix}-mlp-kernel"] = f"{hf_prefix}.mlp.linear_fc1.weight"
    mapping[f"{prefix}-mlp-bias"] = f"{hf_prefix}.mlp.linear_fc1.bias"
    mapping[f"{prefix}-mlp_out-kernel"] = f"{hf_prefix}.mlp.linear_fc2.weight"
    mapping[f"{prefix}-mlp_out-bias"] = f"{hf_prefix}.mlp.linear_fc2.bias"

  # Vision merger_list (deep mergers at layers 8, 16, 24)
  deepstack_indexes = vision_config.get("deepstack_visual_indexes", [8, 16, 24])
  for merger_idx, _ in enumerate(deepstack_indexes):
    prefix = f"params-vision_encoder-Qwen3OmniMoeVisionEncoder_0-merger_{merger_idx}"
    hf_prefix = f"thinker.visual.merger_list.{merger_idx}"

    mapping[f"{prefix}-ln_q-scale"] = f"{hf_prefix}.ln_q.weight"
    mapping[f"{prefix}-ln_q-bias"] = f"{hf_prefix}.ln_q.bias"
    mapping[f"{prefix}-mlp_0-kernel"] = f"{hf_prefix}.mlp.0.weight"
    mapping[f"{prefix}-mlp_0-bias"] = f"{hf_prefix}.mlp.0.bias"
    mapping[f"{prefix}-mlp_2-kernel"] = f"{hf_prefix}.mlp.2.weight"
    mapping[f"{prefix}-mlp_2-bias"] = f"{hf_prefix}.mlp.2.bias"

  # Vision projector (final merger)
  mapping["params-vision_encoder-Qwen3OmniMoeVisionProjector_0-merger-ln_q-scale"] = "thinker.visual.merger.ln_q.weight"
  mapping["params-vision_encoder-Qwen3OmniMoeVisionProjector_0-merger-ln_q-bias"] = "thinker.visual.merger.ln_q.bias"
  mapping["params-vision_encoder-Qwen3OmniMoeVisionProjector_0-merger-mlp_0-kernel"] = (
      "thinker.visual.merger.mlp.0.weight"
  )
  mapping["params-vision_encoder-Qwen3OmniMoeVisionProjector_0-merger-mlp_0-bias"] = "thinker.visual.merger.mlp.0.bias"
  mapping["params-vision_encoder-Qwen3OmniMoeVisionProjector_0-merger-mlp_2-kernel"] = (
      "thinker.visual.merger.mlp.2.weight"
  )
  mapping["params-vision_encoder-Qwen3OmniMoeVisionProjector_0-merger-mlp_2-bias"] = "thinker.visual.merger.mlp.2.bias"

  # Audio mapping
  audio_config = config["thinker_config"]["audio_config"]
  n_audio_layers = audio_config["encoder_layers"]

  # Audio conv layers (3 Conv2D layers for downsampling)
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv2d1-kernel"] = "thinker.audio_tower.conv2d1.weight"
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv2d1-bias"] = "thinker.audio_tower.conv2d1.bias"
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv2d2-kernel"] = "thinker.audio_tower.conv2d2.weight"
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv2d2-bias"] = "thinker.audio_tower.conv2d2.bias"
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv2d3-kernel"] = "thinker.audio_tower.conv2d3.weight"
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv2d3-bias"] = "thinker.audio_tower.conv2d3.bias"

  # Audio conv output projection
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv_out-kernel"] = "thinker.audio_tower.conv_out.weight"

  # Audio encoder layers (32 layers)
  for i in range(n_audio_layers):
    prefix = f"params-audio_encoder-Qwen3OmniAudioEncoder_0-layers_{i}"
    hf_prefix = f"thinker.audio_tower.layers.{i}"

    # Layer norms
    mapping[f"{prefix}-input_layer_norm-scale"] = f"{hf_prefix}.self_attn_layer_norm.weight"
    mapping[f"{prefix}-input_layer_norm-bias"] = f"{hf_prefix}.self_attn_layer_norm.bias"
    mapping[f"{prefix}-post_attention_layer_norm-scale"] = f"{hf_prefix}.final_layer_norm.weight"
    mapping[f"{prefix}-post_attention_layer_norm-bias"] = f"{hf_prefix}.final_layer_norm.bias"

    # Attention (separate Q/K/V)
    mapping[f"{prefix}-self_attention_audio-query-kernel"] = f"{hf_prefix}.self_attn.q_proj.weight"
    mapping[f"{prefix}-self_attention_audio-query-bias"] = f"{hf_prefix}.self_attn.q_proj.bias"
    mapping[f"{prefix}-self_attention_audio-key-kernel"] = f"{hf_prefix}.self_attn.k_proj.weight"
    mapping[f"{prefix}-self_attention_audio-key-bias"] = f"{hf_prefix}.self_attn.k_proj.bias"
    mapping[f"{prefix}-self_attention_audio-value-kernel"] = f"{hf_prefix}.self_attn.v_proj.weight"
    mapping[f"{prefix}-self_attention_audio-value-bias"] = f"{hf_prefix}.self_attn.v_proj.bias"
    mapping[f"{prefix}-self_attention_audio-out-kernel"] = f"{hf_prefix}.self_attn.out_proj.weight"
    mapping[f"{prefix}-self_attention_audio-out-bias"] = f"{hf_prefix}.self_attn.out_proj.bias"

    # MLP (AudioMLP has 2 linear layers: fc1 and fc2)
    mapping[f"{prefix}-AudioMLP-wi-kernel"] = f"{hf_prefix}.fc1.weight"
    mapping[f"{prefix}-AudioMLP-wi-bias"] = f"{hf_prefix}.fc1.bias"
    mapping[f"{prefix}-AudioMLP-wo-kernel"] = f"{hf_prefix}.fc2.weight"
    mapping[f"{prefix}-AudioMLP-wo-bias"] = f"{hf_prefix}.fc2.bias"

  # Audio post layer norm
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-layernorm_post-scale"] = "thinker.audio_tower.ln_post.weight"
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-layernorm_post-bias"] = "thinker.audio_tower.ln_post.bias"

  # Audio projector (2 linear layers)
  mapping["params-audio_encoder-Qwen3OmniAudioProjector_0-proj1-kernel"] = "thinker.audio_tower.proj1.weight"
  mapping["params-audio_encoder-Qwen3OmniAudioProjector_0-proj1-bias"] = "thinker.audio_tower.proj1.bias"
  mapping["params-audio_encoder-Qwen3OmniAudioProjector_0-proj2-kernel"] = "thinker.audio_tower.proj2.weight"
  mapping["params-audio_encoder-Qwen3OmniAudioProjector_0-proj2-bias"] = "thinker.audio_tower.proj2.bias"

  return mapping


def QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
  """Creates parameter transformation functions for Qwen3-Omni.

  This function provides a dictionary of transformation functions (hooks) for
  converting Qwen3-Omni model parameters between MaxText and Hugging Face formats.
  It handles embedding padding and kernel reshaping.

  Args:
    config (dict): Model configuration dictionary, including
      'num_hidden_layers' and optionally 'num_experts'.
    scan_layers (bool, optional): Whether the model uses scanned layers.
      Defaults to False.
    saving_to_hf (bool, optional): The direction of conversion. True for
      MaxText to Hugging Face, False for the reverse. Defaults to False.

  Returns:
    dict: A dictionary mapping MaxText parameter names to their corresponding
      transformation functions.
  """
  # Collect all modality hooks
  mapping = {}

  # Text hooks, reusing QWEN3-MOE hook function
  num_experts_text = config["thinker_config"]["text_config"].get("num_experts", 0)
  n_layers_text = config["thinker_config"]["text_config"]["num_hidden_layers"]
  text_hooks = QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN(
      config={"num_hidden_layers": n_layers_text, "num_experts": num_experts_text},
      maxtext_config=maxtext_config,
      scan_layers=scan_layers,
      saving_to_hf=saving_to_hf,
  )
  mapping.update(text_hooks)

  # Vision hooks
  vision_config = config["thinker_config"]["vision_config"]
  n_vision_layers = vision_config["depth"]
  hidden_size = vision_config["hidden_size"]

  def reshape_kernel_vision(input_tensor, target_shape):
    """Reshape kernel for vision layers."""
    if saving_to_hf:
      flipped_target_shape = np.flip(np.array(target_shape))
      return input_tensor.reshape(flipped_target_shape).T
    else:
      return input_tensor.T.reshape(target_shape)

  def reshape_conv3d_patch_embed(input_tensor, target_shape):
    """Reshape 3D conv patch embedding weight.
    HF: (out_channels, in_channels, temporal, height, width)
    MaxText: (temporal, height, width, in_channels, out_channels)
    """
    if saving_to_hf:
      # MaxText -> HF: (T, H, W, C_in, C_out) -> (C_out, C_in, T, H, W)
      return input_tensor.transpose(4, 3, 0, 1, 2)
    else:
      # HF -> MaxText: (C_out, C_in, T, H, W) -> (T, H, W, C_in, C_out)
      return input_tensor.transpose(2, 3, 4, 1, 0)

  def split_qkv_query(input_tensor, target_shape):
    """Extract Q from fused QKV for HF->MaxText conversion.
    HF has fused QKV: (3*hidden_size, hidden_size)
    MaxText Q: (hidden_size, num_heads, head_dim)
    """
    if saving_to_hf:
      # MaxText -> HF: will be handled by fusion hook
      raise NotImplementedError("Use fusion hook for MaxText->HF")
    else:
      # HF -> MaxText: Extract Q from fused QKV
      # input_tensor shape: (3*hidden_size, hidden_size)
      q_weight = input_tensor[:hidden_size, :]  # (hidden_size, hidden_size)
      return q_weight.T.reshape(target_shape)  # (hidden_size, num_heads, head_dim)

  def split_qkv_key(input_tensor, target_shape):
    """Extract K from fused QKV for HF->MaxText conversion."""
    if saving_to_hf:
      raise NotImplementedError("Use fusion hook for MaxText->HF")
    else:
      # Extract K from fused QKV
      k_weight = input_tensor[hidden_size : 2 * hidden_size, :]
      return k_weight.T.reshape(target_shape)

  def split_qkv_value(input_tensor, target_shape):
    """Extract V from fused QKV for HF->MaxText conversion."""
    if saving_to_hf:
      raise NotImplementedError("Use fusion hook for MaxText->HF")
    else:
      # Extract V from fused QKV
      v_weight = input_tensor[2 * hidden_size :, :]
      return v_weight.T.reshape(target_shape)

  def split_qkv_bias_query(input_tensor, target_shape):
    """Extract Q bias from fused QKV bias."""
    if saving_to_hf:
      raise NotImplementedError("Use fusion hook for MaxText->HF")
    else:
      q_bias = input_tensor[:hidden_size]
      return q_bias.reshape(target_shape)  # (num_heads, head_dim)

  def split_qkv_bias_key(input_tensor, target_shape):
    """Extract K bias from fused QKV bias."""
    if saving_to_hf:
      raise NotImplementedError("Use fusion hook for MaxText->HF")
    else:
      k_bias = input_tensor[hidden_size : 2 * hidden_size]
      return k_bias.reshape(target_shape)

  def split_qkv_bias_value(input_tensor, target_shape):
    """Extract V bias from fused QKV bias."""
    if saving_to_hf:
      raise NotImplementedError("Use fusion hook for MaxText->HF")
    else:
      v_bias = input_tensor[2 * hidden_size :]
      return v_bias.reshape(target_shape)

  def reshape_vision_attn_out(input_tensor, target_shape):
    """Reshape vision attention output projection.
    HF: (hidden_size, hidden_size)
    MaxText: (num_heads, head_dim, hidden_size)
    """
    if saving_to_hf:
      # MaxText -> HF: (num_heads, head_dim, hidden_size) -> (hidden_size, hidden_size)
      return input_tensor.reshape(hidden_size, hidden_size).T
    else:
      # HF -> MaxText: (hidden_size, hidden_size) -> (num_heads, head_dim, hidden_size)
      return input_tensor.T.reshape(target_shape)

  # Vision patch embedding
  mapping["params-vision_encoder-Qwen3OmniMoeVisionEncoder_0-patch_embed-proj-kernel"] = reshape_conv3d_patch_embed

  # Vision blocks
  for i in range(n_vision_layers):
    prefix = f"params-vision_encoder-Qwen3OmniMoeVisionEncoder_0-blocks_{i}"

    # Attention Q/K/V - split from fused QKV
    mapping[f"{prefix}-attn-attn-query-kernel"] = split_qkv_query
    mapping[f"{prefix}-attn-attn-query-bias"] = split_qkv_bias_query
    mapping[f"{prefix}-attn-attn-key-kernel"] = split_qkv_key
    mapping[f"{prefix}-attn-attn-key-bias"] = split_qkv_bias_key
    mapping[f"{prefix}-attn-attn-value-kernel"] = split_qkv_value
    mapping[f"{prefix}-attn-attn-value-bias"] = split_qkv_bias_value

    # Attention output
    mapping[f"{prefix}-attn-attn-out-kernel"] = reshape_vision_attn_out
    # attn-attn-out-bias doesn't need a hook (no reshape needed)

    # MLP
    mapping[f"{prefix}-mlp-kernel"] = reshape_kernel_vision
    mapping[f"{prefix}-mlp_out-kernel"] = reshape_kernel_vision

  # Vision merger_list and projector MLPs
  deepstack_indexes = vision_config.get("deepstack_visual_indexes", [8, 16, 24])
  for merger_idx, _ in enumerate(deepstack_indexes):
    prefix = f"params-vision_encoder-Qwen3OmniMoeVisionEncoder_0-merger_{merger_idx}"
    mapping[f"{prefix}-mlp_0-kernel"] = reshape_kernel_vision
    mapping[f"{prefix}-mlp_2-kernel"] = reshape_kernel_vision

  # Vision projector (final merger)
  mapping["params-vision_encoder-Qwen3OmniMoeVisionProjector_0-merger-mlp_0-kernel"] = reshape_kernel_vision
  mapping["params-vision_encoder-Qwen3OmniMoeVisionProjector_0-merger-mlp_2-kernel"] = reshape_kernel_vision

  # Audio hooks
  audio_config = config["thinker_config"]["audio_config"]
  n_audio_layers = audio_config["encoder_layers"]
  hidden_size_audio = audio_config["d_model"]

  def reshape_kernel_audio(input_tensor, target_shape):
    """Reshape kernel for audio layers."""
    if saving_to_hf:
      flipped_target_shape = np.flip(np.array(target_shape))
      return input_tensor.reshape(flipped_target_shape).T
    else:
      return input_tensor.T.reshape(target_shape)

  def reshape_conv2d_audio(input_tensor, target_shape):
    """Reshape Conv2D weight for audio.

    HF: (out_channels, in_channels, height, width)
    MaxText: (height, width, in_channels, out_channels)
    """
    if saving_to_hf:
      # MaxText -> HF: (H, W, C_in, C_out) -> (C_out, C_in, H, W)
      return input_tensor.transpose(3, 2, 0, 1)
    else:
      # HF -> MaxText: (C_out, C_in, H, W) -> (H, W, C_in, C_out)
      return input_tensor.transpose(2, 3, 1, 0)

  def reshape_audio_attn_qkv(input_tensor, target_shape):
    """Reshape audio attention Q/K/V projection.

    HF: (hidden_size, hidden_size)
    MaxText: (hidden_size, num_heads, head_dim)
    """
    if saving_to_hf:
      # MaxText -> HF: (hidden_size, num_heads, head_dim) -> (hidden_size, hidden_size)
      return input_tensor.reshape(hidden_size_audio, hidden_size_audio).T
    else:
      # HF -> MaxText: (hidden_size, hidden_size) -> (hidden_size, num_heads, head_dim)
      return input_tensor.T.reshape(target_shape)

  def reshape_audio_attn_out(input_tensor, target_shape):
    """Reshape audio attention output projection.
    F
        HF: (hidden_size, hidden_size)
        MaxText: (num_heads, head_dim, hidden_size)
    """
    if saving_to_hf:
      # MaxText -> HF: (num_heads, head_dim, hidden_size) -> (hidden_size, hidden_size)
      return input_tensor.reshape(hidden_size_audio, hidden_size_audio).T
    else:
      # HF -> MaxText: (hidden_size, hidden_size) -> (num_heads, head_dim, hidden_size)
      return input_tensor.T.reshape(target_shape)

  def reshape_audio_attn_qkv_bias(input_tensor, target_shape):
    """Reshape audio attention Q/K/V bias.

    HF: (hidden_size,)
    MaxText: (num_heads, head_dim)
    """
    if saving_to_hf:
      # MaxText -> HF: (num_heads, head_dim) -> (hidden_size,)
      return input_tensor.reshape(hidden_size_audio)
    else:
      # HF -> MaxText: (hidden_size,) -> (num_heads, head_dim)
      return input_tensor.reshape(target_shape)

  # Audio conv layers
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv2d1-kernel"] = reshape_conv2d_audio
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv2d2-kernel"] = reshape_conv2d_audio
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv2d3-kernel"] = reshape_conv2d_audio

  # Audio conv output projection
  mapping["params-audio_encoder-Qwen3OmniAudioEncoder_0-conv_out-kernel"] = reshape_kernel_audio

  # Audio encoder layers
  for i in range(n_audio_layers):
    prefix = f"params-audio_encoder-Qwen3OmniAudioEncoder_0-layers_{i}"

    # Attention Q/K/V
    mapping[f"{prefix}-self_attention_audio-query-kernel"] = reshape_audio_attn_qkv
    mapping[f"{prefix}-self_attention_audio-query-bias"] = reshape_audio_attn_qkv_bias
    mapping[f"{prefix}-self_attention_audio-key-kernel"] = reshape_audio_attn_qkv
    mapping[f"{prefix}-self_attention_audio-key-bias"] = reshape_audio_attn_qkv_bias
    mapping[f"{prefix}-self_attention_audio-value-kernel"] = reshape_audio_attn_qkv
    mapping[f"{prefix}-self_attention_audio-value-bias"] = reshape_audio_attn_qkv_bias

    # Attention output
    mapping[f"{prefix}-self_attention_audio-out-kernel"] = reshape_audio_attn_out

    # MLP
    mapping[f"{prefix}-AudioMLP-wi-kernel"] = reshape_kernel_audio
    mapping[f"{prefix}-AudioMLP-wo-kernel"] = reshape_kernel_audio

  # Audio projector
  mapping["params-audio_encoder-Qwen3OmniAudioProjector_0-proj1-kernel"] = reshape_kernel_audio
  mapping["params-audio_encoder-Qwen3OmniAudioProjector_0-proj2-kernel"] = reshape_kernel_audio

  return mapping


def QWEN3_NNX_TO_VLLM_PARAM_HOOK_FN(target_shape=None):
  """Creates parameter transformation functions for Qwen3.

  This function provides a dictionary of transformation functions (hooks) for
  converting Qwen3 model parameters between NNX and vLLM formats.

  Returns:
    dict: A dictionary mapping NNX parameter names to their corresponding
      transformation functions.
  """
  return {}


def LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
  """
  Returns a dictionary mapping from MaxText parameter names to
  HuggingFace LLaMA3.1 parameter names.

  Args:
      config (dict): Model configuration dictionary containing:
          - num_hidden_layers (int): The number of decoder layers.
      scan_layers (bool, optional): If True, MaxText layers are 'stacked'
          into a single param. Defaults to False.

  Returns:
      dict: A mapping where keys are `atomic_mt_key` (single MaxText parameter names).
        Values are either a single string (unscanned form) or a list of strings
        (scanned form) for stacked layers when `scan_layers=True`.
  """
  n_layers = config["num_hidden_layers"]

  mapping = {
      "params-token_embedder-embedding": "model.embed_tokens.weight",
      "params-decoder-logits_dense-kernel": "lm_head.weight",
      "params-decoder-decoder_norm-scale": "model.norm.weight",
  }

  if scan_layers:
    mapping["params-decoder-layers-self_attention-query-kernel"] = [
        f"model.layers.{layer_idx}.self_attn.q_proj.weight" for layer_idx in range(n_layers)
    ]
    mapping["params-decoder-layers-self_attention-key-kernel"] = [
        f"model.layers.{layer_idx}.self_attn.k_proj.weight" for layer_idx in range(n_layers)
    ]
    mapping["params-decoder-layers-self_attention-value-kernel"] = [
        f"model.layers.{layer_idx}.self_attn.v_proj.weight" for layer_idx in range(n_layers)
    ]
    mapping["params-decoder-layers-self_attention-out-kernel"] = [
        f"model.layers.{layer_idx}.self_attn.o_proj.weight" for layer_idx in range(n_layers)
    ]
    mapping["params-decoder-layers-mlp-wi_0-kernel"] = [
        f"model.layers.{layer_idx}.mlp.gate_proj.weight" for layer_idx in range(n_layers)
    ]
    mapping["params-decoder-layers-mlp-wi_1-kernel"] = [
        f"model.layers.{layer_idx}.mlp.up_proj.weight" for layer_idx in range(n_layers)
    ]
    mapping["params-decoder-layers-mlp-wo-kernel"] = [
        f"model.layers.{layer_idx}.mlp.down_proj.weight" for layer_idx in range(n_layers)
    ]
    mapping["params-decoder-layers-pre_self_attention_layer_norm-scale"] = [
        f"model.layers.{layer_idx}.input_layernorm.weight" for layer_idx in range(n_layers)
    ]
    mapping["params-decoder-layers-post_self_attention_layer_norm-scale"] = [
        f"model.layers.{layer_idx}.post_attention_layernorm.weight" for layer_idx in range(n_layers)
    ]
  else:
    for layer_idx in range(n_layers):
      mapping[f"params-decoder-layers_{layer_idx}-self_attention-query-kernel"] = (
          f"model.layers.{layer_idx}.self_attn.q_proj.weight"
      )
      mapping[f"params-decoder-layers_{layer_idx}-self_attention-key-kernel"] = (
          f"model.layers.{layer_idx}.self_attn.k_proj.weight"
      )
      mapping[f"params-decoder-layers_{layer_idx}-self_attention-value-kernel"] = (
          f"model.layers.{layer_idx}.self_attn.v_proj.weight"
      )
      mapping[f"params-decoder-layers_{layer_idx}-self_attention-out-kernel"] = (
          f"model.layers.{layer_idx}.self_attn.o_proj.weight"
      )
      mapping[f"params-decoder-layers_{layer_idx}-mlp-wi_0-kernel"] = f"model.layers.{layer_idx}.mlp.gate_proj.weight"
      mapping[f"params-decoder-layers_{layer_idx}-mlp-wi_1-kernel"] = f"model.layers.{layer_idx}.mlp.up_proj.weight"
      mapping[f"params-decoder-layers_{layer_idx}-mlp-wo-kernel"] = f"model.layers.{layer_idx}.mlp.down_proj.weight"
      mapping[f"params-decoder-layers_{layer_idx}-pre_self_attention_layer_norm-scale"] = (
          f"model.layers.{layer_idx}.input_layernorm.weight"
      )
      mapping[f"params-decoder-layers_{layer_idx}-post_self_attention_layer_norm-scale"] = (
          f"model.layers.{layer_idx}.post_attention_layernorm.weight"
      )

  return mapping


def LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
  """Creates parameter transformation functions for converting between MaxText and
  HuggingFace formats.

  This function generates a mapping of transformation functions that handle the necessary
  conversions between MaxText and HuggingFace parameter formats, including operations like
  reshaping.
  """
  nlayers = config["num_hidden_layers"]

  def scale_query_layer(input_tensor, target_shape):
    if saving_to_hf:
      depth_scale = np.dtype("float32").type(np.sqrt(config["head_dim"]))
      original_dtype = input_tensor.dtype
      output_tensor = input_tensor.astype(np.float32) * depth_scale
      return output_tensor.astype(original_dtype)
    else:
      depth_scale = np.dtype("float32").type(1 / np.sqrt(config["head_dim"]))
      original_dtype = input_tensor.dtype
      output_tensor = input_tensor.astype(np.float32) * depth_scale
      return output_tensor.astype(original_dtype)

  def adjust_rope(input_tensor, target_shape):
    arr = input_tensor
    if saving_to_hf:
      # Convert from MaxText's interleaved layout to HF's concatenated layout
      evens = arr[..., ::2]
      odds = arr[..., 1::2]
      return jax.numpy.concatenate((evens, odds), axis=arr.ndim - 1)
    else:
      # Convert from HF's concatenated layout to MaxText's interleaved layout
      half_dim = arr.shape[-1] // 2
      first_half = arr[..., :half_dim]
      second_half = arr[..., half_dim:]
      return jax.numpy.stack([first_half, second_half], axis=-1).reshape(arr.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()
    else:
      return input_tensor.transpose().reshape(target_shape)

  # caveat: hook order does affect result
  # to_huggingface
  query_hook_chain = [scale_query_layer, adjust_rope, reshape_kernel]
  key_hook_chain = [adjust_rope, reshape_kernel]
  # to_maxtext
  if not saving_to_hf:
    query_hook_chain.reverse()
    key_hook_chain.reverse()

  hook_fns = {}

  hook_fns["params-decoder-logits_dense-kernel"] = reshape_kernel

  if scan_layers:
    hook_fns = {
        **hook_fns,
        "params-decoder-layers-self_attention-query-kernel": query_hook_chain,
        "params-decoder-layers-self_attention-key-kernel": key_hook_chain,
        "params-decoder-layers-self_attention-value-kernel": reshape_kernel,
        "params-decoder-layers-self_attention-out-kernel": reshape_kernel,
        "params-decoder-layers-mlp-wi_0-kernel": reshape_kernel,
        "params-decoder-layers-mlp-wi_1-kernel": reshape_kernel,
        "params-decoder-layers-mlp-wo-kernel": reshape_kernel,
    }
  else:
    for layer_idx in range(nlayers):
      hook_fns[f"params-decoder-layers_{layer_idx}-self_attention-query-kernel"] = query_hook_chain
      hook_fns[f"params-decoder-layers_{layer_idx}-self_attention-key-kernel"] = key_hook_chain
      hook_fns[f"params-decoder-layers_{layer_idx}-self_attention-value-kernel"] = reshape_kernel
      hook_fns[f"params-decoder-layers_{layer_idx}-self_attention-out-kernel"] = reshape_kernel
      hook_fns[f"params-decoder-layers_{layer_idx}-mlp-wi_0-kernel"] = reshape_kernel
      hook_fns[f"params-decoder-layers_{layer_idx}-mlp-wi_1-kernel"] = reshape_kernel
      hook_fns[f"params-decoder-layers_{layer_idx}-mlp-wo-kernel"] = reshape_kernel
  return hook_fns


def LLAMA31_NNX_TO_VLLM_PARAM_HOOK_FN():
  """Defines and returns hook functions for weight transformations.

  These hooks are applied to specific weights during the conversion
  from MaxText to a HuggingFace-compatible format. They handle
  transformations like RoPE reordering and query scaling that are not
  simple re-mappings.

  Returns:
    A dictionary where keys are MaxText parameter names and values are
    the corresponding transformation functions.
  """

  def reorder_rope(arr):
    """Reorders Rotary Position Embedding (RoPE) weights.

    This function is necessary because MaxText and HuggingFace's vLLM
    implementations may have different orderings for RoPE dimensions.
    It splits the last dimension into even and odd indices and
    concatenates them.

    Args:
      arr: The input weight array.

    Returns:
      The reordered weight array.
    """
    evens = arr[..., ::2]
    odds = arr[..., 1::2]
    return jax.numpy.concatenate((evens, odds), axis=arr.ndim - 1)

  def transform_query_kernel(arr):
    """Transforms the query kernel.

    This involves scaling the kernel by the square root of the head
    dimension and then applying RoPE reordering.

    Args:
      arr: The query kernel weight array.

    Returns:
      The transformed query kernel array.
    """
    head_dim = arr.shape[-1]
    depth_scale = np.dtype("float32").type(np.sqrt(head_dim))
    arr = arr * depth_scale
    return reorder_rope(arr)

  hook_fns = {
      "base.decoder.layers.self_attention.query.kernel": transform_query_kernel,
      "base.decoder.layers.self_attention.key.kernel": reorder_rope,
  }
  return hook_fns


def MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
  """
  Generates the mapping of parameter names from MaxText to Hugging Face for Mixtral.

  Returns:
    dict: A mapping where keys are `atomic_mt_key` (single MaxText parameter names). Values
      are Hugging Face parameter names in one of four forms: unscanned string,
      scanned list of strings, unscanned with expert stacking (list of strings),
      or scanned with expert stacking (nested list of strings).
  """
  mapping = {}

  # Top-level, non-layer-specific parameters
  mapping["params-token_embedder-embedding"] = "model.embed_tokens.weight"
  mapping["params-decoder-decoder_norm-scale"] = "model.norm.weight"
  mapping["params-decoder-logits_dense-kernel"] = "lm_head.weight"

  num_experts = maxtext_config.num_experts

  if scan_layers:
    # Initialize lists for scanned layer weights
    mapping.update(
        {
            "params-decoder-layers-self_attention-query-kernel": [],
            "params-decoder-layers-self_attention-key-kernel": [],
            "params-decoder-layers-self_attention-value-kernel": [],
            "params-decoder-layers-self_attention-out-kernel": [],
            "params-decoder-layers-pre_self_attention_layer_norm-scale": [],
            "params-decoder-layers-post_self_attention_layer_norm-scale": [],
            "params-decoder-layers-MoeBlock_0-gate-kernel": [],
            "params-decoder-layers-MoeBlock_0-wi_0": [],
            "params-decoder-layers-MoeBlock_0-wi_1": [],
            "params-decoder-layers-MoeBlock_0-wo": [],
        }
    )

    for i in range(config["num_hidden_layers"]):
      hf_prefix = f"model.layers.{i}"
      # Attention weights
      mapping["params-decoder-layers-self_attention-query-kernel"].append(f"{hf_prefix}.self_attn.q_proj.weight")
      mapping["params-decoder-layers-self_attention-key-kernel"].append(f"{hf_prefix}.self_attn.k_proj.weight")
      mapping["params-decoder-layers-self_attention-value-kernel"].append(f"{hf_prefix}.self_attn.v_proj.weight")
      mapping["params-decoder-layers-self_attention-out-kernel"].append(f"{hf_prefix}.self_attn.o_proj.weight")

      # RMSNorm weights
      mapping["params-decoder-layers-pre_self_attention_layer_norm-scale"].append(f"{hf_prefix}.input_layernorm.weight")
      mapping["params-decoder-layers-post_self_attention_layer_norm-scale"].append(
          f"{hf_prefix}.post_attention_layernorm.weight"
      )

      # MoE gate
      mapping["params-decoder-layers-MoeBlock_0-gate-kernel"].append(f"{hf_prefix}.block_sparse_moe.gate.weight")

    # Outer loop as experts and inner loop as layers to align with logic in _build_multi_axis_stacked_tensor()
    for j in range(num_experts):
      w1_layers = []
      w3_layers = []
      w2_layers = []

      for i in range(config["num_hidden_layers"]):
        hf_prefix = f"model.layers.{i}"
        w1_layers.append(f"{hf_prefix}.block_sparse_moe.experts.{j}.w1.weight")
        w3_layers.append(f"{hf_prefix}.block_sparse_moe.experts.{j}.w3.weight")
        w2_layers.append(f"{hf_prefix}.block_sparse_moe.experts.{j}.w2.weight")

      mapping["params-decoder-layers-MoeBlock_0-wi_0"].append(w1_layers)
      mapping["params-decoder-layers-MoeBlock_0-wi_1"].append(w3_layers)
      mapping["params-decoder-layers-MoeBlock_0-wo"].append(w2_layers)

  else:
    for i in range(config["num_hidden_layers"]):
      maxtext_prefix = f"params-decoder-layers_{i}"
      hf_prefix = f"model.layers.{i}"

      # Attention weights
      mapping[f"{maxtext_prefix}-self_attention-query-kernel"] = f"{hf_prefix}.self_attn.q_proj.weight"
      mapping[f"{maxtext_prefix}-self_attention-key-kernel"] = f"{hf_prefix}.self_attn.k_proj.weight"
      mapping[f"{maxtext_prefix}-self_attention-value-kernel"] = f"{hf_prefix}.self_attn.v_proj.weight"
      mapping[f"{maxtext_prefix}-self_attention-out-kernel"] = f"{hf_prefix}.self_attn.o_proj.weight"

      # RMSNorm weights
      mapping[f"{maxtext_prefix}-pre_self_attention_layer_norm-scale"] = f"{hf_prefix}.input_layernorm.weight"
      mapping[f"{maxtext_prefix}-post_self_attention_layer_norm-scale"] = f"{hf_prefix}.post_attention_layernorm.weight"

      # MoE gate
      mapping[f"{maxtext_prefix}-MoeBlock_0-gate-kernel"] = f"{hf_prefix}.block_sparse_moe.gate.weight"

      # MoE expert weights (1 MaxText param -> 8 HF params)
      w1_experts = [f"{hf_prefix}.block_sparse_moe.experts.{j}.w1.weight" for j in range(num_experts)]
      w3_experts = [f"{hf_prefix}.block_sparse_moe.experts.{j}.w3.weight" for j in range(num_experts)]
      w2_experts = [f"{hf_prefix}.block_sparse_moe.experts.{j}.w2.weight" for j in range(num_experts)]

      mapping[f"{maxtext_prefix}-MoeBlock_0-wi_0"] = w1_experts
      mapping[f"{maxtext_prefix}-MoeBlock_0-wi_1"] = w3_experts
      mapping[f"{maxtext_prefix}-MoeBlock_0-wo"] = w2_experts

  return mapping


def MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
  """
  Generates parameter conversion hooks for Mixtral between MaxText and Hugging Face.
  """
  hooks = {}

  def reshape_and_transpose_attention(x, target_shape):
    """MaxText: [hidden, n_heads, h_dim] <-> HF: [n_heads * h_dim, hidden]"""
    if saving_to_hf:
      # (H, N, D) -> (H, N*D) -> (N*D, H)
      return x.reshape(config["hidden_size"], -1).transpose()
    else:
      # (N*D, H) -> (H, N*D) -> (H, N, D)
      return x.transpose().reshape(target_shape)

  def reshape_kernel(x, target_shape):
    return x.transpose()

  def scale_query_layer(input_tensor, target_shape):
    if saving_to_hf:
      depth_scale = np.dtype("float32").type(np.sqrt(maxtext_config.head_dim))
      return (input_tensor * depth_scale).astype(input_tensor.dtype)
    else:
      depth_scale = np.dtype("float32").type(1 / np.sqrt(maxtext_config.head_dim))
      return (input_tensor * depth_scale).astype(input_tensor.dtype)

  # hook order does not affect result
  query_hook_chain = [reshape_and_transpose_attention, scale_query_layer]

  if scan_layers:
    plan = [
        ("params-decoder-layers-self_attention-query-kernel", query_hook_chain),
        ("params-decoder-layers-self_attention-key-kernel", reshape_and_transpose_attention),
        ("params-decoder-layers-self_attention-value-kernel", reshape_and_transpose_attention),
        ("params-decoder-layers-self_attention-out-kernel", reshape_and_transpose_attention),
        ("params-decoder-layers-MoeBlock_0-wi_0", reshape_kernel),
        ("params-decoder-layers-MoeBlock_0-wi_1", reshape_kernel),
        ("params-decoder-layers-MoeBlock_0-wo", reshape_kernel),
        ("params-decoder-layers-MoeBlock_0-gate-kernel", reshape_kernel),
    ]
  else:
    plan = [
        ("params-decoder-layers_{i}-self_attention-query-kernel", query_hook_chain),
        ("params-decoder-layers_{i}-self_attention-key-kernel", reshape_and_transpose_attention),
        ("params-decoder-layers_{i}-self_attention-value-kernel", reshape_and_transpose_attention),
        ("params-decoder-layers_{i}-self_attention-out-kernel", reshape_and_transpose_attention),
        ("params-decoder-layers_{i}-MoeBlock_0-wi_0", reshape_kernel),
        ("params-decoder-layers_{i}-MoeBlock_0-wi_1", reshape_kernel),
        ("params-decoder-layers_{i}-MoeBlock_0-wo", reshape_kernel),
        ("params-decoder-layers_{i}-MoeBlock_0-gate-kernel", reshape_kernel),
    ]
  plan.append(("params-decoder-logits_dense-kernel", reshape_kernel))

  for maxtext_pattern, op_func in plan:
    if "{i}" in maxtext_pattern:
      for i in range(config["num_hidden_layers"]):
        hooks[maxtext_pattern.format(i=i)] = op_func
    else:
      hooks[maxtext_pattern] = op_func
  return hooks


def OLMO3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
  """Returns mapping from MaxText to HuggingFace Olmo3 weight paths."""

  # Olmo3 uses an inhomogeneous layer cycle (4 layers: 3 sliding, 1 global).
  # MaxText handles this by defining sub-layers (layers_0, layers_1...) within a block.
  n_layers = config["num_hidden_layers"]

  # Default Olmo3 cycle length is 4 if not specified in config
  layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval

  # Base mapping for embeddings and global norms
  mapping = {
      "params-token_embedder-embedding": "model.embed_tokens.weight",
      "params-decoder-decoder_norm-scale": "model.norm.weight",
      "params-decoder-logits_dense-kernel": "lm_head.weight",
  }

  if scan_layers:
    # Scanned: Map MaxText 'layers_k' to HF layers [k, k+cycle, k+2*cycle, ...]
    for cycle_idx in range(layer_cycle_interval):
      hf_indices = range(cycle_idx, n_layers, layer_cycle_interval)
      prefix = f"params-decoder-layers-layers_{cycle_idx}"

      mapping.update(
          {
              # Attention Projections
              f"{prefix}-attention-query-kernel": [f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices],
              f"{prefix}-attention-key-kernel": [f"model.layers.{i}.self_attn.k_proj.weight" for i in hf_indices],
              f"{prefix}-attention-value-kernel": [f"model.layers.{i}.self_attn.v_proj.weight" for i in hf_indices],
              f"{prefix}-attention-out-kernel": [f"model.layers.{i}.self_attn.o_proj.weight" for i in hf_indices],
              # QK Norms (Olmo3 Specific)
              f"{prefix}-attention-query_norm-scale": [f"model.layers.{i}.self_attn.q_norm.weight" for i in hf_indices],
              f"{prefix}-attention-key_norm-scale": [f"model.layers.{i}.self_attn.k_norm.weight" for i in hf_indices],
              # MLP (wi_0=gate, wi_1=up, wo=down)
              f"{prefix}-mlp-wi_0-kernel": [f"model.layers.{i}.mlp.gate_proj.weight" for i in hf_indices],
              f"{prefix}-mlp-wi_1-kernel": [f"model.layers.{i}.mlp.up_proj.weight" for i in hf_indices],
              f"{prefix}-mlp-wo-kernel": [f"model.layers.{i}.mlp.down_proj.weight" for i in hf_indices],
              # Layer Norms
              f"{prefix}-post_self_attention_layer_norm-scale": [
                  f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices
              ],
              f"{prefix}-post_mlp_layer_norm-scale": [
                  f"model.layers.{i}.post_feedforward_layernorm.weight" for i in hf_indices
              ],
          }
      )

  else:
    # Unscanned: Direct 1-to-1 mapping
    for i in range(n_layers):
      prefix = f"params-decoder-layers_{i}"
      hf_prefix = f"model.layers.{i}"

      mapping.update(
          {
              f"{prefix}-attention-query-kernel": f"{hf_prefix}.self_attn.q_proj.weight",
              f"{prefix}-attention-key-kernel": f"{hf_prefix}.self_attn.k_proj.weight",
              f"{prefix}-attention-value-kernel": f"{hf_prefix}.self_attn.v_proj.weight",
              f"{prefix}-attention-out-kernel": f"{hf_prefix}.self_attn.o_proj.weight",
              f"{prefix}-attention-query_norm-scale": f"{hf_prefix}.self_attn.q_norm.weight",
              f"{prefix}-attention-key_norm-scale": f"{hf_prefix}.self_attn.k_norm.weight",
              f"{prefix}-mlp-wi_0-kernel": f"{hf_prefix}.mlp.gate_proj.weight",
              f"{prefix}-mlp-wi_1-kernel": f"{hf_prefix}.mlp.up_proj.weight",
              f"{prefix}-mlp-wo-kernel": f"{hf_prefix}.mlp.down_proj.weight",
              f"{prefix}-post_self_attention_layer_norm-scale": f"{hf_prefix}.post_attention_layernorm.weight",
              f"{prefix}-post_mlp_layer_norm-scale": f"{hf_prefix}.post_feedforward_layernorm.weight",
          }
      )

  return mapping


def OLMO3_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
  """Creates parameter transformation functions for Olmo3."""

  # 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
    else:
      return input_tensor.T.reshape(target_shape)

  # Identity mapping for Norms
  # Olmo3 checkpoints typically have weights ~1.0.
  # If MaxText RMSNorm adds 1.0 (x * (1+w)), we might need w-1.0.
  # However, if weights are zeroed/mismatched, identity is safer to restore logic flow.
  def scale_rmsnorm_layer(input_tensor, target_shape):
    return input_tensor.reshape(target_shape)

  # Identity mapping for QK Norms (assuming MaxText attentions.py was patched to use global norm)
  def adapt_olmo3_qk_norm(input_tensor, target_shape):
    return input_tensor.reshape(target_shape)

  # Padding for Embedding layer (Vocab size adjustments)
  def pad_hf_embedding_layer(input_tensor, target_shape):
    source_vocab_size = input_tensor.shape[0]
    target_vocab_size = target_shape[0]
    if source_vocab_size == target_vocab_size:
      return input_tensor
    if saving_to_hf:
      return input_tensor[:target_vocab_size, :]
    else:
      padded_tensor = np.zeros(target_shape, dtype=input_tensor.dtype)
      padded_tensor[:source_vocab_size, :] = input_tensor
      return padded_tensor

  hooks = {
      "params-token_embedder-embedding": pad_hf_embedding_layer,
      "params-decoder-logits_dense-kernel": reshape_kernel,
      "params-decoder-decoder_norm-scale": scale_rmsnorm_layer,
  }

  kernel_keys = [
      "attention-query-kernel",
      "attention-key-kernel",
      "attention-value-kernel",
      "attention-out-kernel",
      "mlp-wi_0-kernel",
      "mlp-wi_1-kernel",
      "mlp-wo-kernel",
  ]

  norm_keys = [
      "post_self_attention_layer_norm-scale",
      "post_mlp_layer_norm-scale",
      "attention-query_norm-scale",
      "attention-key_norm-scale",
  ]

  cycle_len = getattr(maxtext_config, "inhomogeneous_layer_cycle_interval", 4)
  n_layers = config["num_hidden_layers"]

  if scan_layers:
    for cycle_idx in range(cycle_len):
      prefix = f"params-decoder-layers-layers_{cycle_idx}"
      for key in kernel_keys:
        hooks[f"{prefix}-{key}"] = reshape_kernel
      for key in norm_keys:
        # For QK norm, we use the specific adaptor (which is currently Identity
        # but separates logic if we need to revert to averaging later)
        if "attention-" in key and "_norm-scale" in key:
          hooks[f"{prefix}-{key}"] = adapt_olmo3_qk_norm
        else:
          hooks[f"{prefix}-{key}"] = scale_rmsnorm_layer
  else:
    for i in range(n_layers):
      prefix = f"params-decoder-layers_{i}"
      for key in kernel_keys:
        hooks[f"{prefix}-{key}"] = reshape_kernel
      for key in norm_keys:
        if "attention-" in key and "_norm-scale" in key:
          hooks[f"{prefix}-{key}"] = adapt_olmo3_qk_norm
        else:
          hooks[f"{prefix}-{key}"] = scale_rmsnorm_layer

  return hooks


# {maxtext model name: {maxtext weight name: hf weight name}}
PARAM_MAPPING = {
    "gemma2-2b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING,
    "gemma2-9b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING,
    "gemma2-27b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING,
    "gemma3-4b": GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING,
    "gemma3-12b": GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING,
    "gemma3-27b": GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen2.5-7b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen2.5-14b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-0.6b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-1.7b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-1.7b-base": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-4b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-4b-base": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-4b-thinking-2507": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-8b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-8b-base": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-14b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-14b-base": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-32b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING,
    "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING,
    "llama3.1-405b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-30b-a3b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-30b-a3b-base": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-235b-a22b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-coder-480b-a35b": QWEN_MAXTEXT_TO_HF_PARAM_MAPPING,
    "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING,
    "gpt-oss-20b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING,
    "gpt-oss-120b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING,
    "qwen3-next-80b-a3b": QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING,
    "mixtral-8x7b": MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING,
    "mixtral-8x22b": MIXTRAL_MAXTEXT_TO_HF_PARAM_MAPPING,
    "olmo3-7b": OLMO3_MAXTEXT_TO_HF_PARAM_MAPPING,
    "olmo3-7b-pt": OLMO3_MAXTEXT_TO_HF_PARAM_MAPPING,
    "olmo3-32b": OLMO3_MAXTEXT_TO_HF_PARAM_MAPPING,
}

# {maxtext model name: {maxtext weight name: bi-directional transform}}
HOOK_FNS = {
    "gemma2-2b": GEMMA2_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "gemma2-9b": GEMMA2_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "gemma2-27b": GEMMA2_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "gemma3-4b": GEMMA3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "gemma3-12b": GEMMA3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "gemma3-27b": GEMMA3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen2.5-7b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen2.5-14b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-0.6b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-1.7b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-1.7b-base": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-4b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-4b-base": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-4b-thinking-2507": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-8b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-8b-base": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-14b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-14b-base": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-32b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "llama3.1-405b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-30b-a3b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-30b-a3b-base": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-235b-a22b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-coder-480b-a35b": QWEN_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "gpt-oss-20b": GPT_OSS_TO_HF_PARAM_HOOK_FN,
    "gpt-oss-120b": GPT_OSS_TO_HF_PARAM_HOOK_FN,
    "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "qwen3-next-80b-a3b": QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "mixtral-8x7b": MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "mixtral-8x22b": MIXTRAL_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "olmo3-7b": OLMO3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "olmo3-7b-pt": OLMO3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
    "olmo3-32b": OLMO3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
}

VLLM_HOOK_FNS = {
    "qwen3": QWEN3_NNX_TO_VLLM_PARAM_HOOK_FN,
    "llama3.1": LLAMA31_NNX_TO_VLLM_PARAM_HOOK_FN,
    "deepseek3": DEEPSEEK_NNX_TO_VLLM_PARAM_HOOK_FN,
}