def build_llamalike_block(
    name: str,
    init_base_rng: jax.Array | None,
    config: LlamalikeTransformerConfig,
    block_index: int | None = None,
) -> model_parts.TransformerBlock:
  """Builds a transformer block from a configuration.

  Args:
    name: Name of the block.
    init_base_rng: Base RNG for initializing the parameters.
    config: The configuration of the model.
    block_index: The index of the transformer block in the list of blocks. Can
      be None if the attention type doesn't depend on the block index.

  Returns:
    A full transformer block.
  """
  attn_sequence = [
      pz.nn.RMSLayerNorm.from_config(
          name=f"{name}/pre_attention_norm",
          init_base_rng=init_base_rng,
          across_axes={"embedding": config.embedding_dim},
          dtype=config.parameter_dtype,
          epsilon=config.rms_norm_eps,
      ),
      build_llamalike_attention(
          f"{name}/attention",
          init_base_rng,
          config,
          block_index=block_index,
      ),
  ]
  if config.use_post_attn_norm:
    attn_sequence.append(
        pz.nn.RMSLayerNorm.from_config(
            name=f"{name}/post_attention_norm",
            init_base_rng=init_base_rng,
            across_axes={"embedding": config.embedding_dim},
            dtype=config.parameter_dtype,
            epsilon=config.rms_norm_eps,
        )
    )
  ffw_sequence = [
      pz.nn.RMSLayerNorm.from_config(
          name=f"{name}/pre_ffw_norm",
          init_base_rng=init_base_rng,
          across_axes={"embedding": config.embedding_dim},
          dtype=config.parameter_dtype,
          epsilon=config.rms_norm_eps,
      ),
      build_llamalike_feedforward(f"{name}/mlp", init_base_rng, config),
  ]
  if config.use_post_ffw_norm:
    ffw_sequence.append(
        pz.nn.RMSLayerNorm.from_config(
            name=f"{name}/post_ffw_norm",
            init_base_rng=init_base_rng,
            across_axes={"embedding": config.embedding_dim},
            dtype=config.parameter_dtype,
            epsilon=config.rms_norm_eps,
        )
    )
  return model_parts.TransformerBlock(
      sublayers=[
          pz.nn.Residual(pz.nn.Sequential(attn_sequence)),
          pz.nn.Residual(pz.nn.Sequential(ffw_sequence)),
      ],
  )
Builds a transformer block from a configuration.

  Args:
    name: Name of the block.
    init_base_rng: Base RNG for initializing the parameters.
    config: The configuration of the model.
    block_index: The index of the transformer block in the list of blocks. Can
      be None if the attention type doesn't depend on the block index.

  Returns:
    A full transformer block.
  
penzai/models/transformer/variants/llamalike_common.py