def gpt_neox_from_huggingface_model(
    model: GPTNeoXForCausalLM,
    upcast_activations_to_float32: bool = False,
    use_layer_stack: bool = False,
) -> model_parts.TransformerLM:
  """Converts a GPT-NeoX model to a Penzai model.

  This function converts GPT-NeoX models from their HuggingFace implementations
  to Penzai.

  Note: Checkpoint conversion is only implemented for the most common set of
  hyperparameters for GPT-NeoX models, including GPT-NeoX-20B and the Pythia
  scaling suite.

  Args:
    model: The HuggingFace GPT-NeoX model.
    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.
  """
  try:
    import transformers  # pylint: disable=import-outside-toplevel
  except ImportError as exc:
    raise RuntimeError("HuggingFace transformers is not available") from exc

  if type(model) is not transformers.GPTNeoXForCausalLM:  # pylint: disable=unidiomatic-typecheck
    raise ValueError(
        "gpt_neox_from_huggingface_model should be called with a"
        f" GPTNeoXForCausalLM instance, but got {type(model).__name__}."
    )

  hf_config = model.config
  # Check any modified configuration arguments against the base config to make
  # sure we support newer architecture features. (Assumes that new features are
  # added in a backwards-compatible way and do not change the defaults for the
  # configuration class.)
  hf_config_attributes = hf_config.to_dict()
  reference_attributes = transformers.GPTNeoXConfig().to_dict()
  handled_or_ignored_attributes = {
      # Handled during conversion:
      "hidden_act",
      "hidden_size",
      "intermediate_size",
      "layer_norm_eps",
      "num_attention_heads",
      "num_hidden_layers",
      "rotary_emb_base",
      "rotary_pct",
      "vocab_size",
      # Ignored by conversion:
      "max_position_embeddings",
      "torch_dtype",
      "architectures",
      "bos_token_id",
      "eos_token_id",
      "_attn_implementation_autoset",
      "head_dim",
  }
  bad_attributes = {}
  for k, v in hf_config_attributes.items():
    if k in handled_or_ignored_attributes or (
        k in reference_attributes and v == reference_attributes[k]
    ):
      pass
    else:
      bad_attributes[k] = v
  if bad_attributes:
    raise ValueError(
        "Conversion of a GPTNeoXForCausalLM does not support these"
        f" configuration attributes: {repr(bad_attributes)}"
    )

  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

  projection_dim = hf_config.hidden_size // hf_config.num_attention_heads
  converted_activation_fn = {
      "gelu": "gelu_exact",
      "gelu_new": "gelu_approx",
      "relu": "relu",
      "silu": "silu",
  }[hf_config.hidden_act]
  pz_config = GPTNeoXTransformerConfig(
      num_attention_heads=hf_config.num_attention_heads,
      embedding_dim=hf_config.hidden_size,
      projection_dim=projection_dim,
      mlp_hidden_dim=hf_config.intermediate_size,
      num_decoder_blocks=hf_config.num_hidden_layers,
      vocab_size=hf_config.vocab_size,
      activation_fn=converted_activation_fn,
      rope_wavelength=hf_config.rotary_emb_base,
      rope_subset_size=int(hf_config.rotary_pct * projection_dim),
      layernorm_epsilon=hf_config.layer_norm_eps,
      parameter_dtype=param_dtype,
      activation_dtype=activation_dtype,
      use_layer_stack=use_layer_stack,
  )
  model_def = build_gpt_neox_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["gpt_neox.embed_in.weight"]
      ).tag("vocabulary", "embedding"),
      "final_norm/scale.weights": pz.nx.NamedArray.wrap(
          converted["gpt_neox.final_layer_norm.weight"]
      ).tag("embedding"),
      "final_norm/shift.bias": pz.nx.NamedArray.wrap(
          converted["gpt_neox.final_layer_norm.bias"]
      ).tag("embedding"),
      "lm_head.weights": pz.nx.NamedArray.wrap(
          converted["embed_out.weight"]
      ).tag("vocabulary", "embedding"),
  }

  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"gpt_neox.layers.{i}.input_layernorm.weight"]
        ).tag("embedding")
    )
    cur_block_params["pre_attention_norm/shift.bias"] = pz.nx.NamedArray.wrap(
        converted[f"gpt_neox.layers.{i}.input_layernorm.bias"]
    ).tag("embedding")
    cur_block_params["pre_ffw_norm/scale.weights"] = pz.nx.NamedArray.wrap(
        converted[f"gpt_neox.layers.{i}.post_attention_layernorm.weight"]
    ).tag("embedding")
    cur_block_params["pre_ffw_norm/shift.bias"] = pz.nx.NamedArray.wrap(
        converted[f"gpt_neox.layers.{i}.post_attention_layernorm.bias"]
    ).tag("embedding")
    cur_block_params["mlp/in/Linear.weights"] = pz.nx.NamedArray.wrap(
        converted[f"gpt_neox.layers.{i}.mlp.dense_h_to_4h.weight"]
    ).tag("neurons", "embedding")
    cur_block_params["mlp/in/AddBias.bias"] = pz.nx.NamedArray.wrap(
        converted[f"gpt_neox.layers.{i}.mlp.dense_h_to_4h.bias"]
    ).tag("neurons")
    cur_block_params["mlp/out/Linear.weights"] = pz.nx.NamedArray.wrap(
        converted[f"gpt_neox.layers.{i}.mlp.dense_4h_to_h.weight"]
    ).tag("embedding", "neurons")
    cur_block_params["mlp/out/AddBias.bias"] = pz.nx.NamedArray.wrap(
        converted[f"gpt_neox.layers.{i}.mlp.dense_4h_to_h.bias"]
    ).tag("embedding")

    qkv_weights = (
        pz.nx.NamedArray.wrap(
            converted[f"gpt_neox.layers.{i}.attention.query_key_value.weight"]
        )
        .reshape((
            pz_config.num_attention_heads,
            3,
            pz_config.projection_dim,
            pz_config.embedding_dim,
        ))
        .tag("heads", "qkv", "projection", "embedding")
    )
    qkv_biases = (
        pz.nx.NamedArray.wrap(
            converted[f"gpt_neox.layers.{i}.attention.query_key_value.bias"]
        )
        .reshape((
            pz_config.num_attention_heads,
            3,
            pz_config.projection_dim,
        ))
        .tag("heads", "qkv", "projection")
    )

    cur_block_params["attention/query/Linear.weights"] = qkv_weights[{"qkv": 0}]
    cur_block_params["attention/key/Linear.weights"] = qkv_weights[{"qkv": 1}]
    cur_block_params["attention/value/Linear.weights"] = qkv_weights[{"qkv": 2}]
    cur_block_params["attention/query/AddBias.bias"] = qkv_biases[{"qkv": 0}]
    cur_block_params["attention/key/AddBias.bias"] = qkv_biases[{"qkv": 1}]
    cur_block_params["attention/value/AddBias.bias"] = qkv_biases[{"qkv": 2}]

    cur_block_params["attention/output/Linear.weights"] = (
        pz.nx.NamedArray.wrap(
            converted[f"gpt_neox.layers.{i}.attention.dense.weight"]
        )
        .reshape((
            pz_config.embedding_dim,
            pz_config.num_attention_heads,
            pz_config.projection_dim,
        ))
        .tag("embedding", "heads", "projection")
    )
    cur_block_params["attention/output/AddBias.bias"] = pz.nx.NamedArray.wrap(
        converted[f"gpt_neox.layers.{i}.attention.dense.bias"]
    ).tag("embedding")

  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 GPT-NeoX model to a Penzai model.

  This function converts GPT-NeoX models from their HuggingFace implementations
  to Penzai.

  Note: Checkpoint conversion is only implemented for the most common set of
  hyperparameters for GPT-NeoX models, including GPT-NeoX-20B and the Pythia
  scaling suite.

  Args:
    model: The HuggingFace GPT-NeoX model.
    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.
  
penzai/models/transformer/variants/gpt_neox.py