def build_llamalike_attention(
    name: str,
    init_base_rng: jax.Array | None,
    config: LlamalikeTransformerConfig,
    block_index: int | None = None,
) -> pz.nn.Attention:
  """Builds an attention block from a configuration.

  Args:
    name: Name of the attention 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:
    An Attention block.
  """
  embedding_dim = config.embedding_dim
  projection_dim = config.projection_dim

  common_head_axes, qkv_einsum, query_only_head_axes, q_einsum = _head_info(
      config
  )

  if config.query_scaling_factor == "default":
    query_scaling_factor = projection_dim**-0.5
  else:
    query_scaling_factor = config.query_scaling_factor

  # As used in https://github.com/google-deepmind/gemma.
  # (This exact value is probably not important.)
  masked_out_value = jnp.array(-2.3819763e38, dtype=config.activation_dtype)

  if isinstance(config.attention_type, AttentionType):
    attention_type = config.attention_type
  else:
    if block_index is None:
      raise ValueError(
          "block_index must be specified if attention_type is a sequence."
      )
    attention_type = config.attention_type[
        block_index % len(config.attention_type)
    ]

  if isinstance(attention_type, AttentionTypeSlidingWindowCausal):
    attn_masker = pz.nn.ApplyCausalSlidingWindowAttentionMask(
        sliding_window_size=attention_type.window_size,
        masked_out_value=masked_out_value,
    )
  elif isinstance(attention_type, AttentionTypeGlobalCausal):
    attn_masker = pz.nn.ApplyCausalAttentionMask(
        masked_out_value=masked_out_value,
    )
  else:
    raise ValueError(f"Unsupported attention type {attention_type}")

  query_key_to_attn_sublayers = [
      pz.nn.NamedEinsum(
          (
              {"seq": "tq", **qkv_einsum, **q_einsum, "projection": "p"},
              {"seq": "tkv", **qkv_einsum, "projection": "p"},
          ),
          {"seq": "tq", **qkv_einsum, **q_einsum, "kv_seq": "tkv"},
      ),
  ]
  if config.attn_logits_soft_cap is not None:
    query_key_to_attn_sublayers.append(
        pz.nn.TanhSoftCap(
            soft_cap=jnp.array(
                config.attn_logits_soft_cap, dtype=config.activation_dtype
            )
        )
    )
  query_key_to_attn_sublayers.extend([
      attn_masker,
      pz.nn.Softmax("kv_seq"),
  ])

  return pz.nn.Attention(
      input_to_query=pz.nn.Sequential([
          pz.nn.Linear.from_config(
              name=f"{name}/query",
              init_base_rng=init_base_rng,
              input_axes={"embedding": embedding_dim},
              output_axes={
                  **common_head_axes,
                  **query_only_head_axes,
                  "projection": projection_dim,
              },
              dtype=config.parameter_dtype,
          ),
          pz.nn.ApplyRoPE(
              positions_input_name="token_positions",
              embedding_axis="projection",
              max_wavelength=config.rope_wavelength,
          ),
          pz.nn.ConstantRescale(
              by=jnp.array(query_scaling_factor, dtype=config.activation_dtype)
          ),
      ]),
      input_to_key=pz.nn.Sequential([
          pz.nn.Linear.from_config(
              name=f"{name}/key",
              init_base_rng=init_base_rng,
              input_axes={"embedding": embedding_dim},
              output_axes={**common_head_axes, "projection": projection_dim},
              dtype=config.parameter_dtype,
          ),
          pz.nn.ApplyRoPE(
              positions_input_name="token_positions",
              embedding_axis="projection",
              max_wavelength=config.rope_wavelength,
          ),
      ]),
      input_to_value=pz.nn.Sequential([
          pz.nn.Linear.from_config(
              name=f"{name}/value",
              init_base_rng=init_base_rng,
              input_axes={"embedding": embedding_dim},
              output_axes={**common_head_axes, "projection": projection_dim},
              dtype=config.parameter_dtype,
          ),
      ]),
      query_key_to_attn=pz.nn.Sequential(query_key_to_attn_sublayers),
      attn_value_to_output=pz.nn.Sequential([
          pz.nn.NamedEinsum(
              (
                  {"seq": "tq", **qkv_einsum, **q_einsum, "kv_seq": "tkv"},
                  {"seq": "tkv", **qkv_einsum, "projection": "p"},
              ),
              {"seq": "tq", **qkv_einsum, **q_einsum, "projection": "p"},
          ),
          pz.nn.Linear.from_config(
              name=f"{name}/output",
              init_base_rng=init_base_rng,
              input_axes={
                  **common_head_axes,
                  **query_only_head_axes,
                  "projection": projection_dim,
              },
              output_axes={"embedding": embedding_dim},
              dtype=config.parameter_dtype,
          ),
      ]),
  )
Builds an attention block from a configuration.

  Args:
    name: Name of the attention 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:
    An Attention block.
  
penzai/models/transformer/variants/llamalike_common.py