def from_config(cls, config: GemmaTransformerConfig) -> GemmaTransformer:
    """Constructs an uninitialized transformer with the Gemma architecture.

    Args:
      config: The configuration of the Gemma model.

    Returns:
      A ``GemmaTransformer`` with uninitialized parameters. All side input
      effects will have already been appropriately handled.
    """
    # Embedding table is shared between first and last layers.
    emb_table = pz.nn.mark_shareable(
        pz.nn.add_parameter_prefix(
            "embedder",
            pz.nn.EmbeddingTable.from_config(
                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))
    sublayers.append(
        pz.nn.ConstantRescale(
            by=jnp.sqrt(config.embedding_dim).astype(config.activation_dtype)
        )
    )
    for i in range(config.num_decoder_blocks):
      sublayers.append(
          pz.nn.add_parameter_prefix(
              f"block_{i}", GemmaTransformerBlock.from_config(config)
          )
      )
    sublayers.extend([
        pz.nn.add_parameter_prefix(
            "final_norm",
            pz.nn.RMSLayerNorm.from_config(
                across_axes={"embedding": config.embedding_dim},
                dtype=config.parameter_dtype,
            ),
        ),
        pz.nn.EmbeddingDecode(emb_table),
    ])
    # Handle shared parameters and side inputs.
    return GemmaTransformer(
        config=config,
        body=pz.de.WithSideInputsFromInputTuple.handling(
            pz.nn.attach_shared_parameters(pz.nn.Sequential(sublayers)),
            tags=["token_positions", "attn_mask"],
        ),
    )
Constructs an uninitialized transformer with the Gemma architecture.

    Args:
      config: The configuration of the Gemma model.

    Returns:
      A ``GemmaTransformer`` with uninitialized parameters. All side input
      effects will have already been appropriately handled.
    
penzai/deprecated/v1/example_models/gemma/model_core.py