def build_gpt_neox_transformer(
    config: GPTNeoXTransformerConfig,
    init_base_rng: jax.Array | None = None,
    name: str = "transformer",
) -> model_parts.TransformerLM:
  """Builds a Llama-like transformer model from a configuration.

  Args:
    config: The configuration of the model.
    init_base_rng: Base RNG for initializing the parameters.
    name: Name for the top-level model, used as a prefix for all parameters.

  Returns:
    A full transformer model.
  """

  # Embedding table is shared between first and last layers.
  emb_table = pz.nn.EmbeddingTable.from_config(
      name=f"{name}/embedder",
      init_base_rng=init_base_rng,
      vocab_size=config.vocab_size,
      embedding_axes={"embedding": config.embedding_dim},
      dtype=config.parameter_dtype,
  )
  sublayers = []
  sublayers.append(pz.nn.EmbeddingLookup(emb_table))
  if config.activation_dtype != config.parameter_dtype:
    sublayers.append(pz.nn.CastToDType(config.activation_dtype))

  if config.use_layer_stack:
    sublayers.append(
        pz.nn.LayerStack.from_sublayer_builder(
            builder=build_gpt_neox_block,
            stack_axis="blocks",
            stack_axis_size=config.num_decoder_blocks,
            init_base_rng=init_base_rng,
            builder_kwargs=dict(name=f"{name}/blocks", config=config),
        )
    )
  else:
    for i in range(config.num_decoder_blocks):
      sublayers.append(
          build_gpt_neox_block(f"{name}/block_{i}", init_base_rng, config)
      )

  sublayers.append(
      pz.nn.LayerNorm.from_config(
          name=f"{name}/final_norm",
          init_base_rng=init_base_rng,
          across_axes={"embedding": config.embedding_dim},
          dtype=config.parameter_dtype,
      )
  )

  sublayers.append(
      pz.nn.Linear.from_config(
          name=f"{name}/lm_head",
          init_base_rng=init_base_rng,
          input_axes={"embedding": config.embedding_dim},
          output_axes={"vocabulary": config.vocab_size},
      )
  )

  return model_parts.TransformerLM(
      metadata=model_parts.TransformerMetadata(
          common_head_axes={"heads": config.num_attention_heads},
          query_only_head_axes={},
          embedding_dim=config.embedding_dim,
          projection_dim=config.projection_dim,
          mlp_hidden_dim=config.mlp_hidden_dim,
          vocab_size=config.vocab_size,
          activation_dtype=config.activation_dtype,
      ),
      body=pz.nn.Sequential(sublayers),
  )
Builds a Llama-like transformer model from a configuration.

  Args:
    config: The configuration of the model.
    init_base_rng: Base RNG for initializing the parameters.
    name: Name for the top-level model, used as a prefix for all parameters.

  Returns:
    A full transformer model.
  
penzai/models/transformer/variants/gpt_neox.py