def from_pretrained(
      cls,
      ckpt_params: dict[str, Any],
      upcast_activations_to_float32: bool = False,
  ) -> GemmaTransformer:
    """Constructs a ``GemmaTransformer`` from an official Gemma Flax checkpoint.

    The parameters of the loaded ``GemmaTransformer`` will be close to those in
    the original checkpoint with a few modifications:

    * Query, key, and value heads are stored in three separate matrices instead
      of being stored either as a single matrix (qkv_einsum) or as two (q_einsum
      and kv_einsum).

    * `RMSLayerNorm` weights have their values increased by one, instead of
      adding one at call time.

    * Axes of parameters are identified by name instead of by position.

    Args:
      ckpt_params: Nested dictionary of weights from the Gemma checkpoint.
      upcast_activations_to_float32: Whether to cast activations to float32 when
        the model runs. This is useful for doing interpretability research at
        higher precision without consuming additional memory.

    Returns:
      A GemmaTransformer model containing the loaded parameters.
    """
    params = {k.removeprefix("transformer/"): v for k, v in ckpt_params.items()}
    num_layers = 0
    for i in itertools.count():
      if f"layer_{i}/mlp/linear" not in params:
        num_layers = i
        break
    hidden_dim, embed_dim = params["layer_0/mlp/linear"]["w"].shape
    attn_0_einsum_param = params["layer_0/attn/attn_vec_einsum"]["w"]
    num_heads, proj_dim, _ = attn_0_einsum_param.shape
    single_kv_head = "layer_0/attn/qkv_einsum" not in params
    vocab_size = params["embedder"]["input_embedding"].shape[0]

    if upcast_activations_to_float32:
      activation_dtype = jnp.float32
    else:
      activation_dtype = attn_0_einsum_param.dtype

    config = GemmaTransformerConfig(
        num_heads=num_heads,
        embedding_dim=embed_dim,
        projection_dim=proj_dim,
        single_kv_head=single_kv_head,
        mlp_hidden_dim=hidden_dim,
        num_decoder_blocks=num_layers,
        vocab_size=vocab_size,
        parameter_dtype=attn_0_einsum_param.dtype,
        activation_dtype=activation_dtype,
    )
    model_def = cls.from_config(config)
    parameter_mapping = {
        "embedder.embeddings": pz.nx.NamedArray.wrap(
            params["embedder"]["input_embedding"]
        ).tag("vocabulary", "embedding"),
        "final_norm.scale.weights": pz.nx.NamedArray.wrap(
            1 + params["final_norm"]["scale"]
        ).tag("embedding"),
    }
    for i in range(config.num_decoder_blocks):
      parameter_mapping.update({
          f"block_{i}.pre_attention_norm.scale.weights": pz.nx.NamedArray.wrap(
              1 + params[f"layer_{i}/pre_attention_norm"]["scale"]
          ).tag("embedding"),
          f"block_{i}.pre_ffw_norm.scale.weights": pz.nx.NamedArray.wrap(
              1 + params[f"layer_{i}/pre_ffw_norm"]["scale"]
          ).tag("embedding"),
          f"block_{i}.mlp.gating_linear.weights": pz.nx.NamedArray.wrap(
              params[f"layer_{i}/mlp/gating_einsum"]["w"][0]
          ).tag("embedding", "neurons"),
          f"block_{i}.mlp.value_linear.weights": pz.nx.NamedArray.wrap(
              params[f"layer_{i}/mlp/gating_einsum"]["w"][1]
          ).tag("embedding", "neurons"),
          f"block_{i}.mlp.out_linear.weights": pz.nx.NamedArray.wrap(
              params[f"layer_{i}/mlp/linear"]["w"]
          ).tag("neurons", "embedding"),
          f"block_{i}.attn.output.weights": pz.nx.NamedArray.wrap(
              params[f"layer_{i}/attn/attn_vec_einsum"]["w"]
          ).tag("heads", "projection", "embedding"),
      })
      if config.single_kv_head:
        parameter_mapping.update({
            f"block_{i}.attn.query.weights": pz.nx.NamedArray.wrap(
                params[f"layer_{i}/attn/q_einsum"]["w"]
            ).tag("heads", "embedding", "projection"),
            f"block_{i}.attn.key.weights": pz.nx.NamedArray.wrap(
                params[f"layer_{i}/attn/kv_einsum"]["w"][0].squeeze(0)
            ).tag("embedding", "projection"),
            f"block_{i}.attn.value.weights": pz.nx.NamedArray.wrap(
                params[f"layer_{i}/attn/kv_einsum"]["w"][1].squeeze(0)
            ).tag("embedding", "projection"),
        })
      else:
        parameter_mapping.update({
            f"block_{i}.attn.query.weights": pz.nx.NamedArray.wrap(
                params[f"layer_{i}/attn/qkv_einsum"]["w"][0]
            ).tag("heads", "embedding", "projection"),
            f"block_{i}.attn.key.weights": pz.nx.NamedArray.wrap(
                params[f"layer_{i}/attn/qkv_einsum"]["w"][1]
            ).tag("heads", "embedding", "projection"),
            f"block_{i}.attn.value.weights": pz.nx.NamedArray.wrap(
                params[f"layer_{i}/attn/qkv_einsum"]["w"][2]
            ).tag("heads", "embedding", "projection"),
        })
    return (
        model_def.select()
        .at_instances_of(pz.nn.UninitializedParameter)
        .apply(
            lambda param: param.initialize_with_value(
                parameter_mapping[param.name],
                strict_dtype=False,
            )
        )
    )
Constructs a ``GemmaTransformer`` from an official Gemma Flax checkpoint.

    The parameters of the loaded ``GemmaTransformer`` will be close to those in
    the original checkpoint with a few modifications:

    * Query, key, and value heads are stored in three separate matrices instead
      of being stored either as a single matrix (qkv_einsum) or as two (q_einsum
      and kv_einsum).

    * `RMSLayerNorm` weights have their values increased by one, instead of
      adding one at call time.

    * Axes of parameters are identified by name instead of by position.

    Args:
      ckpt_params: Nested dictionary of weights from the Gemma checkpoint.
      upcast_activations_to_float32: Whether to cast activations to float32 when
        the model runs. This is useful for doing interpretability research at
        higher precision without consuming additional memory.

    Returns:
      A GemmaTransformer model containing the loaded parameters.
    
penzai/deprecated/v1/example_models/gemma/model_core.py