def llamalike_from_huggingface_model(
    model: Any,
    upcast_activations_to_float32: bool = False,
    use_layer_stack: bool = False,
) -> model_parts.TransformerLM:
  """Converts a "llama-like" HuggingFace model to a Penzai model.

  This function converts Llama-like models from their HuggingFace
  implementations to Penzai. It does not do any checks and blindly assumes
  that the architecture follows the defaults from the Llama model family.
  You may want to use the model-specific wrappers in `variants.llama` or
  `variants.mistral` instead.

  Args:
    model: The HuggingFace model, which is assumed to be similar to the Llama or
      Mistral architectures. (Not all configuration arguments are checked, so
      this may end up producing different behavior if given an incompatible
      configuration.)
    upcast_activations_to_float32: Whether to cast activations to float32 when
      the model runs. This allows analyzing activations at higher precision
      without consuming additional memory for parameters.
    use_layer_stack: Whether to use a layer stack for the decoder blocks.

  Returns:
    A Transformer model containing the loaded parameters, assuming a Llama-like
    architecture.
  """
  hf_config = model.config
  num_kv_heads = hf_config.num_key_value_heads
  query_head_multiplier = hf_config.num_attention_heads // num_kv_heads
  assert num_kv_heads * query_head_multiplier == hf_config.num_attention_heads

  sliding_window_size = getattr(hf_config, "sliding_window", None)
  if sliding_window_size is None:
    attention_type = AttentionTypeGlobalCausal()
  else:
    attention_type = AttentionTypeSlidingWindowCausal(sliding_window_size)

  param_dtype = {
      "torch.float32": jnp.float32,
      "torch.bfloat16": jnp.bfloat16,
  }[str(model.dtype)]
  if upcast_activations_to_float32:
    activation_dtype = jnp.float32
  else:
    activation_dtype = param_dtype

  pz_config = LlamalikeTransformerConfig(
      num_kv_heads=num_kv_heads,
      query_head_multiplier=query_head_multiplier,
      embedding_dim=hf_config.hidden_size,
      projection_dim=hf_config.hidden_size // hf_config.num_attention_heads,
      mlp_hidden_dim=hf_config.intermediate_size,
      num_decoder_blocks=hf_config.num_hidden_layers,
      vocab_size=hf_config.vocab_size,
      mlp_variant="swiglu",
      rope_wavelength=hf_config.rope_theta,
      tie_embedder_and_logits=False,
      attention_type=attention_type,
      rms_norm_eps=hf_config.rms_norm_eps,
      parameter_dtype=param_dtype,
      activation_dtype=activation_dtype,
      use_layer_stack=use_layer_stack,
  )
  model_def = build_llamalike_transformer(
      pz_config, init_base_rng=None, name="transformer"
  )

  state_dict = model.state_dict()
  converted = {k: jax.dlpack.from_dlpack(v) for k, v in state_dict.items()}

  parameter_mapping = {
      "embedder.embeddings": pz.nx.NamedArray.wrap(
          converted["model.embed_tokens.weight"]
      ).tag("vocabulary", "embedding"),
      "final_norm/scale.weights": pz.nx.NamedArray.wrap(
          converted["model.norm.weight"]
      ).tag("embedding"),
      "lm_head.weights": pz.nx.NamedArray.wrap(converted["lm_head.weight"]).tag(
          "vocabulary", "embedding"
      ),
  }

  def fix_qkvo(which, arr):
    arr = pz.nx.wrap(arr)
    if which == "q":
      if pz_config.query_head_multiplier == 1:
        return arr.reshape((
            pz_config.num_kv_heads,
            pz_config.projection_dim,
            pz_config.embedding_dim,
        )).tag("heads", "projection", "embedding")
      elif pz_config.num_kv_heads == 1:
        return arr.reshape((
            pz_config.query_head_multiplier,
            pz_config.projection_dim,
            pz_config.embedding_dim,
        )).tag("query_heads", "projection", "embedding")
      else:
        return arr.reshape((
            pz_config.num_kv_heads,
            pz_config.query_head_multiplier,
            pz_config.projection_dim,
            pz_config.embedding_dim,
        )).tag("head_groups", "query_heads", "projection", "embedding")
    elif which == "k" or which == "v":
      if pz_config.query_head_multiplier == 1:
        return arr.reshape((
            pz_config.num_kv_heads,
            pz_config.projection_dim,
            pz_config.embedding_dim,
        )).tag("heads", "projection", "embedding")
      elif pz_config.num_kv_heads == 1:
        return arr.reshape((
            pz_config.projection_dim,
            pz_config.embedding_dim,
        )).tag("projection", "embedding")
      else:
        return arr.reshape((
            pz_config.num_kv_heads,
            pz_config.projection_dim,
            pz_config.embedding_dim,
        )).tag("head_groups", "projection", "embedding")
    elif which == "o":
      if pz_config.query_head_multiplier == 1:
        return arr.reshape((
            pz_config.embedding_dim,
            pz_config.num_kv_heads,
            pz_config.projection_dim,
        )).tag("embedding", "heads", "projection")
      elif pz_config.num_kv_heads == 1:
        return arr.reshape((
            pz_config.embedding_dim,
            pz_config.query_head_multiplier,
            pz_config.projection_dim,
        )).tag("embedding", "query_heads", "projection")
      else:
        return arr.reshape((
            pz_config.embedding_dim,
            pz_config.num_kv_heads,
            pz_config.query_head_multiplier,
            pz_config.projection_dim,
        )).tag("embedding", "head_groups", "query_heads", "projection")
    else:
      raise NotImplementedError(which)

  all_block_params = []

  for i in range(pz_config.num_decoder_blocks):
    cur_block_params = {}
    all_block_params.append(cur_block_params)

    cur_block_params["pre_attention_norm/scale.weights"] = (
        pz.nx.NamedArray.wrap(
            converted[f"model.layers.{i}.input_layernorm.weight"]
        ).tag("embedding")
    )
    cur_block_params["pre_ffw_norm/scale.weights"] = pz.nx.NamedArray.wrap(
        converted[f"model.layers.{i}.post_attention_layernorm.weight"]
    ).tag("embedding")
    cur_block_params["mlp/gating_linear.weights"] = pz.nx.NamedArray.wrap(
        converted[f"model.layers.{i}.mlp.gate_proj.weight"]
    ).tag("neurons", "embedding")
    cur_block_params["mlp/value_linear.weights"] = pz.nx.NamedArray.wrap(
        converted[f"model.layers.{i}.mlp.up_proj.weight"]
    ).tag("neurons", "embedding")
    cur_block_params["mlp/out_linear.weights"] = pz.nx.NamedArray.wrap(
        converted[f"model.layers.{i}.mlp.down_proj.weight"]
    ).tag("embedding", "neurons")

    cur_block_params["attention/query.weights"] = fix_qkvo(
        "q", converted[f"model.layers.{i}.self_attn.q_proj.weight"]
    )
    cur_block_params["attention/key.weights"] = fix_qkvo(
        "k", converted[f"model.layers.{i}.self_attn.k_proj.weight"]
    )
    cur_block_params["attention/value.weights"] = fix_qkvo(
        "v", converted[f"model.layers.{i}.self_attn.v_proj.weight"]
    )
    cur_block_params["attention/output.weights"] = fix_qkvo(
        "o", converted[f"model.layers.{i}.self_attn.o_proj.weight"]
    )

  if use_layer_stack:
    for key in all_block_params[0].keys():
      vals = [
          all_block_params[i][key] for i in range(pz_config.num_decoder_blocks)
      ]
      parameter_mapping[f"blocks/{key}"] = pz.nx.stack(vals, "blocks")
  else:
    for i in range(pz_config.num_decoder_blocks):
      for key, value in all_block_params[i].items():
        parameter_mapping[f"block_{i}/{key}"] = value

  # Create parameter objects for each parameter.
  model = pz.bind_variables(
      model_def,
      [
          pz.Parameter(value=v, label=f"transformer/{k}")
          for k, v in parameter_mapping.items()
      ],
  )
  pz.nn.assert_no_parameter_slots(model)
  return model
Converts a "llama-like" HuggingFace model to a Penzai model.

  This function converts Llama-like models from their HuggingFace
  implementations to Penzai. It does not do any checks and blindly assumes
  that the architecture follows the defaults from the Llama model family.
  You may want to use the model-specific wrappers in `variants.llama` or
  `variants.mistral` instead.

  Args:
    model: The HuggingFace model, which is assumed to be similar to the Llama or
      Mistral architectures. (Not all configuration arguments are checked, so
      this may end up producing different behavior if given an incompatible
      configuration.)
    upcast_activations_to_float32: Whether to cast activations to float32 when
      the model runs. This allows analyzing activations at higher precision
      without consuming additional memory for parameters.
    use_layer_stack: Whether to use a layer stack for the decoder blocks.

  Returns:
    A Transformer model containing the loaded parameters, assuming a Llama-like
    architecture.
  
penzai/models/transformer/variants/llamalike_common.py