def build_llamalike_transformer(
    config: LlamalikeTransformerConfig,
    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.tie_embedder_and_logits:
    sublayers.append(
        pz.nn.ConstantRescale(
            by=jnp.sqrt(config.embedding_dim).astype(config.activation_dtype)
        )
    )

  if config.use_layer_stack:
    if not isinstance(config.attention_type, AttentionType):
      raise ValueError(
          "Layer stack does not currently support per-layer attention types."
      )
    sublayers.append(
        pz.nn.LayerStack.from_sublayer_builder(
            builder=build_llamalike_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:
    if not isinstance(config.attention_type, AttentionType):
      if config.num_decoder_blocks % len(config.attention_type) != 0:
        raise ValueError(
            "Per-layer attention types must have a length that divides the"
            " number of blocks."
        )
    for block_index in range(config.num_decoder_blocks):
      sublayers.append(
          build_llamalike_block(
              f"{name}/block_{block_index}", init_base_rng, config, block_index
          )
      )

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

  if config.tie_embedder_and_logits:
    sublayers.append(pz.nn.EmbeddingDecode(emb_table))
  else:
    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},
        )
    )

  if config.final_logit_softcap:
    sublayers.append(
        pz.nn.TanhSoftCap(
            soft_cap=jnp.array(
                config.final_logit_softcap, dtype=config.activation_dtype
            )
        )
    )

  common_head_axes, _, query_only_head_axes, _ = _head_info(config)
  return model_parts.TransformerLM(
      metadata=model_parts.TransformerMetadata(
          common_head_axes=common_head_axes,
          query_only_head_axes=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/llamalike_common.py