### Prepare and Batch Tokenized Examples

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Encodes various text examples into token IDs, pads them to a fixed length, and batches them for model input. Requires a vocabulary object and a list of extra examples.

```python
all_toks = []
for extra_example in extra_examples:
  subtoks = [vocab.bos_id()] + vocab.EncodeAsIds(extra_example)
  subtoks = subtoks + [vocab.pad_id()] * (40 - len(subtoks))
  all_toks.append(subtoks[:40])

new_example_batch = pz.nx.wrap(
    jnp.array(all_toks).astype(jnp.int32)
).tag("batch", "seq")
```

--------------------------------

### Install Penzai

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/named_axes.ipynb

Installs the Penzai library with notebook support. Use this in Colab or Kaggle environments if Penzai is not already installed.

```python
try:
  import penzai
except ImportError:
  !pip install penzai[notebook]
```

--------------------------------

### Install Penzai

Source: https://github.com/google-deepmind/penzai/blob/main/README.md

Install the Penzai library using pip. This is a prerequisite for using Penzai.

```bash
pip install penzai
```

--------------------------------

### Create a Basic Selection

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/selectors.ipynb

Demonstrates the starting point for Penzai selections: creating a selection object that initially contains only the root object.

```python
pz.select(my_nested_object)
```

--------------------------------

### Enable Basic Interactive Treescope Setup

Source: https://github.com/google-deepmind/penzai/blob/main/docs/api/treescope.rst

Use this function to enable Treescope with basic interactive setup, particularly for rendering custom types configured with the old __penzai_repr__ extension method.

```python
pz.ts.basic_interactive_setup()
```

--------------------------------

### Conditional Gemma model loading setup

Source: https://github.com/google-deepmind/penzai/blob/main/docs/_include/_glue_figures.ipynb

Sets up logic to load the Gemma model, checking for kagglehub availability and TPU devices. Defaults to loading Gemma if available.

```python
try:
  import kagglehub
  if jax.devices()[0].platform == "tpu":
    load_gemma = True
  else:
    load_gemma = False
  load_gemma = True
except ImportError:
  kagglehub = None
  load_gemma = False
```

--------------------------------

### Set up a Simple Penzai Training Loop

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

Configures and runs a basic training loop for a Penzai model using `basic_training.StatefulTrainer`. Requires example inputs and targets, and a loss function.

```python
from penzai.toolshed import basic_training
import optax
```

```python
example_inputs = pz.nx.wrap(
    jax.random.normal(jax.random.key(100), (100, 8))
).tag("batch", "features")
example_targets = pz.nx.wrap(
    jax.random.normal(jax.random.key(101), (100, 8))
).tag("batch", "features")

def loss_fn(model, rng, state, current_input, current_target):
  del rng, state  # More complex training loops could use these if needed
  model_out = model(current_input)
  losses = pz.nx.nmap(jnp.square)(model_out - current_target)
  loss = losses.untag("batch", "features").unwrap().sum()
  return (loss, None, {"my_loss": loss})

trainer = basic_training.StatefulTrainer.build(
    root_rng=jax.random.key(42),
    model=model,
    optimizer_def=optax.adam(0.01),
    loss_fn=loss_fn,
)

outputs = []
while trainer.state.value.step < 1000:
  out = trainer.step(
      current_input=example_inputs,
      current_target=example_targets,
  )
  if trainer.state.value.step % 20 == 0:
    print(f"At {trainer.state.value.step}: {out}")
```

--------------------------------

### Tokenize Example Text

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Encodes the example text into a JAX array of token IDs using the provided vocabulary. Includes a beginning-of-sequence token.

```python
tokens = jnp.array([vocab.bos_id()] + vocab.EncodeAsIds(example_text))
```

--------------------------------

### Example Usage of KnockOutAttentionHeads

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Demonstrates how to instantiate and use the `KnockOutAttentionHeads` layer to modify attention patterns. This example shows knocking out every other head and visualizing the resulting attention scores.

```python
# Knock out every other head
knockout_layer = KnockOutAttentionHeads(
    head_mask=pz.nx.wrap(jnp.array(
        [1,0,1,0,1,0,1,0,1,0]
    ).astype(jnp.bfloat16)).tag("best_heads")
)

# Show the results (which should alternate kept and knocked-out)
treescope.render_array(
    knockout_layer(top_attn_patterns),
    vmax=1,
    rows=["seq"],
    valid_mask=simple_causal_mask,
)
```

--------------------------------

### Tokenize Batch of Examples

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Converts a list of string examples into padded token IDs suitable for model input. It handles BOS, EOS, and padding tokens.

```python
def tokenize_batch(examples, pad_length=32, include_eos=True):
  padded_tokens = []
  for example in examples:
    example_tokens = [vocab.bos_id()] + vocab.EncodeAsIds(example)
    if include_eos:
      example_tokens = example_tokens + [vocab.eos_id()]
    assert len(example_tokens) <= pad_length
    # Pad from the right (simplifies input positional embeddings)
    example_tokens = (
        example_tokens
        + [vocab.pad_id()] * (pad_length - len(example_tokens))
    )
    padded_tokens.append(example_tokens)
  return pz.nx.wrap(jnp.array(padded_tokens)).tag("batch", "seq")
```

--------------------------------

### Generating Additional Example Sequences

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Generates a list of additional example sequences for testing. These sequences are created by randomly choosing digits and concatenating them with themselves.

```python
extra_examples = []

# Some other random digit sequences:
for i in range(100, 104):
  s = "".join(str(i) for i in jax.random.choice(jax.random.key(i), jnp.arange(10), (20,)))
  extra_examples.append(s + s)
```

--------------------------------

### Generate Synthetic Example

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Creates a single synthetic data example by randomly choosing two numbers, evaluating the mystery function, and formatting the input and output as a string.

```python
def generate_example(np_rng):
  a, b = np_rng.choice(1000, size=(2,)) 
  c = mystery_function(a, b)
  return f">>> mystery_function({a}, {b})\n{c}"
```

--------------------------------

### Instantiating and Using a Penzai Softmax Layer

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

This example shows how to create an instance of the `Softmax` layer defined previously and then apply it to a named array. This demonstrates the practical usage of Penzai layers with named axis interfaces.

```python
layer = Softmax("vocabulary")
layer
```

```python
layer(named)
```

--------------------------------

### Visualize Validation Examples

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Generates and visualizes a batch of synthetic examples used for validation, employing treescope for clear representation of tokenized data.

```python
%%autovisualize treescope.ArrayAutovisualizer.for_tokenizer(vocab)
np_rng = np.random.default_rng(98765)
validation_examples = tokenize_batch([generate_example(np_rng) for _ in range(32)])
validation_examples
```

--------------------------------

### Initialize Rewired Model

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Sets up the initial model checkpoint for rewired experiments. This ensures a consistent starting point for testing different hypotheses about model behavior.

```python
# Start with the original unmodified model checkpoint.
rewired_model = model
```

--------------------------------

### Prepare Example Input Text

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Creates a string of repeated numeric digits to serve as input for the Gemma model. This is inspired by Olsson et al. (2021) to provide a more natural input distribution than random tokens.

```python
example_text = (
    "01976954310149754605"
    + "01976954310149754605"
)
```

--------------------------------

### Build Custom Model using Sequential Combinator

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/howto_reference.md

Implement a custom model by composing Penzai primitives with combinators like `pz.nn.Sequential`. This example shows how to initialize an `Affine` layer within a sequential model, passing along the initialization key and configuring layer dimensions.

```python
def build_my_model(
    name: str,
    init_base_rng: jax.Array | None,
    # ... any other arguments ...
):
  # Initialize all model components and return them, e.g.:
  return pz.nn.Sequential([
      pz.nn.Affine.from_config(
          # Extend the name:
          name=f"{name}/Affine_0",
          # Pass along the initialization key (no need to split it)
          init_base_rng=init_base_rng,
          # Configure the layer (for example)
          input_axes={"features": 8},
          output_axes={"features": 8},
      ),
      # ... Add more layers as needed ...
  ])
```

--------------------------------

### Setup Treescope for Interactive Use

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/named_axes.ipynb

Enables Treescope as the default IPython pretty-printer and activates automatic array visualization. Recommended for interactive Penzai development.

```python
treescope.basic_interactive_setup(autovisualize_arrays=True)
```

--------------------------------

### Advanced Treescope Registration and Autovisualization

Source: https://github.com/google-deepmind/penzai/blob/main/docs/api/treescope.rst

Provides more control over Treescope setup by registering it as the default, enabling autovisualize magic, and setting a global ArrayAutovisualizer. Useful for custom type rendering with legacy Penzai APIs.

```python
pz.ts.register_as_default()
pz.ts.register_autovisualize_magic()
pz.ts.active_autovisualizer.set_globally(pz.ts.ArrayAutovisualizer())
```

--------------------------------

### Display Local JAX Devices and Assert Count

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/jitting_and_sharding.ipynb

Shows the available JAX devices and asserts that the number of local devices is equal to the expected count (e.g., 8). This confirms the environment setup for multi-device operations.

```python
pz.show(jax.local_devices())
assert jax.local_device_count() == 8
```

--------------------------------

### Define Prompts for Sampling

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Defines a list of string prompts to be used for model sampling. Includes examples of direct prompts and a prompt to let the model generate its own problem.

```python
prompts = [
    ">>> mystery_function(123, 123)",
    ">>> mystery_function(101, 15)",
    ">>> mystery_function(999, 876)",
    ">>>", # Let the model write and solve its own problem
]
```

--------------------------------

### Define Counterfactual Examples

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Sets up a list of counterfactual input sequences for testing model behavior under different conditions. Includes sequences for original, nonrepeating, and patched attention patterns.

```python
counterfactuals = [
    "01976954310149754605" + "01976954310149754605",
    "67717010284911166217" + "06302739717444079179",
    "67717010284911166217" + "06302739717444079179",  # <- Patching the non-repeating sequence this time.
]
all_toks = []
for cf_example in counterfactuals:
  subtoks = [vocab.bos_id()] + vocab.EncodeAsIds(cf_example)
  all_toks.append(subtoks)

counterfactuals_batch = pz.nx.wrap(
    jnp.array(all_toks).astype(jnp.int32)
).tag("worlds", "seq")  # <- Name it using the same "worlds" axis convention.
```

--------------------------------

### Download Gemma Model Weights using KaggleHub

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/howto_reference.md

Download pretrained Gemma model weights from Kaggle. This snippet shows how to use `kagglehub.model_download` to get the weights and specify the checkpoint path.

```python
import kagglehub
import orbax.checkpoint
from penzai.models.transformer import variants

# Download Gemma 1 7B:
weights_dir = kagglehub.model_download('google/gemma/Flax/7b')
ckpt_path = os.path.join(weights_dir, '7b')
```

--------------------------------

### Define a Custom Linear Layer in Penzai

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/howto_reference.md

Implement a custom `pz.nn.Layer` for a linear transformation. This example shows how to define parameters, handle non-PyTree fields, implement the `__call__` method, and create a `from_config` builder for parameter initialization.

```python
import dataclasses
import jax
import jax.numpy as jnp

import penzai as pz

@pz.pytree_dataclass
class SimpleLinear(pz.nn.Layer):
  # Parameters are annotated as `ParameterLike` to allow swapping them out after
  # initialization.
  kernel: pz.nn.ParameterLike[pz.nx.NamedArray]

  # Non-Pytree fields (which are not arraylike) should be annotated as such to
  # tell JAX not to try to convert them:
  features_axis: str = dataclasses.field(metadata={"pytree_node": False})

  def __call__(
    self,
    x: pz.nx.NamedArray, /, **unused_side_inputs
  ) -> pz.nx.NamedArray:
    pos_x = x.untag(self.features_axis)
    pos_kernel = self.kernel.value.untag("out_features", "in_features")
    pos_y = pz.nx.nmap(jnp.dot)(pos_kernel, pos_x)
    return pos_y.tag(self.features_axis)

  @classmethod
  def from_config(
      cls,
      name: str,
      init_base_rng: jax.Array | None,
      in_features: int,
      out_features: int,
      features_axis: str = "features",
  ) -> "SimpleLinear":
    """Constructs a linear layer from configuration arguments."""
    def _initializer(key):
      arr = jax.nn.initializers.xavier_normal()(
          key, (out_features, in_features)
      )
      return pz.nx.wrap(arr).tag("out_features", "in_features")

    return cls(
        kernel=pz.nn.make_parameter(
            name=f"{name}.kernel",
            init_base_rng=init_base_rng,
            initializer=_initializer,
        ),
        features_axis=features_axis,
    )
```

--------------------------------

### JIT Compile Parameter Optimization with Shardings

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/jitting_and_sharding.ipynb

Demonstrates how to JIT-compile a function that initializes model parameters, specifying the desired output shardings.

```python
def functional_init(init_base_rng):
  model = ...
  return pz.unbind_variables(model, freeze=True)

sharded_init = jax.jit(
  functional_init,
  out_shardings=..., # <- insert your desired sharding specification here
)

model = pz.bind_variables(*sharded_init(rng))
```

--------------------------------

### Initialize and Visualize a Simple MLP

Source: https://github.com/google-deepmind/penzai/blob/main/README.md

Initialize a simple Multi-Layer Perceptron (MLP) model using Penzai and visualize it. This demonstrates basic model instantiation and output in an interactive environment.

```python
from penzai.models import simple_mlp
mlp = simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(0),
    feature_sizes=[8, 32, 32, 8]
)

# Models and arrays are visualized automatically when you output them from a
# Colab/IPython notebook cell:
mlp
```

--------------------------------

### MLP Configuration with from_config

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

Demonstrates configuring a simple MLP model using its `from_config` class method. This method takes configuration arguments to set up sublayers and parameters.

```python
mlp = simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.PRNGKey(1),
    feature_sizes=[8, 32, 32, 8],
    activation_fn=jax.nn.gelu,
)
mlp
```

--------------------------------

### Get Model Logits

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Passes the prepared token sequence through the Gemma model to obtain the output log-probabilities.

```python
logits = model(token_seq)
```

--------------------------------

### Create and Display a Simple MLP Model

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

Instantiates a simple Multi-Layer Perceptron (MLP) using Penzai's MLP class from configuration. The model is then displayed, leveraging Treescope for visualization.

```python
mlp = simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(0),
    feature_sizes=[8, 32, 32, 8]
)
lp
```

--------------------------------

### Get log probability of correct next token

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Indexes into the vocabulary axis using the correct tokens to retrieve the specific log probability.

```python
log_prob_of_correct_next = sliced_preds[{"vocabulary": correct_next_token}]
```

--------------------------------

### Initialize a Stateful Trainer

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/howto_reference.md

Instantiate `StatefulTrainer` from `penzai.toolshed.basic_training` to set up a basic training loop that automatically updates `Parameter` variables.

```python
stateful_trainer = pz.toolshed.basic_training.StatefulTrainer(model, loss_fn)
```

--------------------------------

### Reshaping PositionalSharding

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/jitting_and_sharding.ipynb

Reshapes a `PositionalSharding` to have multiple axes, allowing for more complex data distribution strategies. This example reshapes a 1D sharding into a 2D sharding.

```python
pos_sharding.reshape((4,2))
```

--------------------------------

### Gradient Computation with `select_and_set_by_path`

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/selectors.ipynb

An alternative method for selective gradient computation using `select_and_set_by_path`. This infers the selection from the provided paths, simplifying the gradient setup.

```python
def my_loss(obj):
  return obj["a"] + jnp.sum(obj["b"]**2) + jnp.sum(obj["c"][0]["value"])

my_float_object = jax.tree_util.tree_map(lambda leaf: jnp.array(leaf, dtype=jnp.float32), my_nested_object)

# Take gradients w.r.t. non-scalars only, ignoring my_float_object["a"] and my_float_object["c"][2]["value"]
# No need to store the gradient selection itself, since it can be inferred from vectors_by_path.
vectors_by_path = pz.select(my_float_object).at_instances_of(jax.Array).where(lambda arr: arr.size > 1).get_by_path()

jax.grad(
    lambda vectors_by_path: my_loss(pz.select(my_float_object).select_and_set_by_path(vectors_by_path))
)(vectors_by_path)
```

--------------------------------

### Initialize a Simple MLP

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Initializes a basic Multi-Layer Perceptron (MLP) with specified feature sizes using Penzai's `from_config` method. Requires JAX and Penzai libraries.

```python
mlp = simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(0),
    feature_sizes=[8, 32, 32, 8],
)
```

--------------------------------

### Initialize Orbax Checkpointer and Get Metadata

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads_2B.ipynb

Initializes an Orbax checkpointer and retrieves metadata for the model checkpoint. This metadata is used for restoring the model parameters.

```python
import orbax.checkpoint

checkpointer = orbax.checkpoint.PyTreeCheckpointer()
metadata = checkpointer.metadata(ckpt_path)
```

--------------------------------

### Model Initialization: V1 vs. V2 Parameter Initialization

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/v2_differences.md

Compares the old method of initializing model parameters using `pz.nn.initialize_parameters` with the new V2 approach, which uses `init_base_rng` during model configuration.

```python
# Old
pz.nn.initialize_parameters(
  simple_mlp.MLP.from_config(feature_sizes=[2, 32, 32, 2]),
  jax.random.key(10),
)

# New
mlp = simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(10),
    feature_sizes=[2, 32, 32, 2],
)
```

--------------------------------

### Retrieve Selected Values by Path

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/selectors.ipynb

Get selected values organized by their paths within the original structure using .get_by_path(). This is equivalent to accessing the .selected_by_path attribute.

```python
selection.get_by_path()
```

--------------------------------

### Visualize Rewiring Instance and Path Matrix

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads_2B.ipynb

Visualizes the instantiated `RewireComputationPaths` layer and its corresponding path matrix.

```python
pz.show(read_both_from_indirect)

```

```python
pz.show(read_both_from_indirect.path_matrix())
```

--------------------------------

### Initialize a Simple MLP

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Creates a simple Multi-Layer Perceptron (MLP) model using Penzai's MLP class. This is a foundational step before applying LoRA.

```python
mlp = simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(0),
    feature_sizes=[2, 32, 32, 2],
)
```

--------------------------------

### Get pretty keystr for Parameters

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/selectors.ipynb

Generates a human-readable string representation of Penzai Parameter keys using `pz.pretty_keystr`. This is helpful for debugging and understanding complex key paths.

```python
[
  pz.pretty_keystr(key, mlp) for key in pz.select(mlp).at_instances_of(pz.Parameter).selected_by_path.keys()
]
```

--------------------------------

### Get JAX tree_util keystr for Parameters

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/selectors.ipynb

Converts Penzai Parameter keys into their string representation using `jax.tree_util.keystr`. This provides a standardized way to view key paths.

```python
[
  jax.tree_util.keystr(key) for key in pz.select(mlp).at_instances_of(pz.Parameter).selected_by_path.keys()
]
```

--------------------------------

### Build Stateful Trainer

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Initializes a Penzai StatefulTrainer for fine-tuning the LoRA-ified Gemma model with AdamW optimizer.

```python
trainer = basic_training.StatefulTrainer.build(
    model=loraified_gemma_model,
    optimizer_def=optax.adamw(5e-5, weight_decay=0.01),
    root_rng=jax.random.key(42),
    loss_fn=xent_loss_fn,
    donate_states=True,
)
```

--------------------------------

### Get Penzai Parameter keys

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/selectors.ipynb

Retrieves a list of all Penzai Parameter keys within a given Penzai object. This is useful for inspecting the structure and identifying trainable parameters.

```python
[
  key for key in pz.select(mlp).at_instances_of(pz.Parameter).selected_by_path.keys()
]
```

--------------------------------

### Import Penzai Core and Treescope

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/named_axes.ipynb

Imports the core Penzai library (aliased as pz) and the Treescope pretty-printer.

```python
import treescope
import penzai
from penzai import pz
```

--------------------------------

### Prepare Counterfactual Sequences

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Creates three sequences for activation patching: a reference sequence, a counterfactual non-repeating sequence, and the reference sequence again for patching.

```python
counterfactuals = [
    "01976954310149754605" + "01976954310149754605",  # <- Our running example so far. This one will be a reference.
    "67717010284911166217" + "06302739717444079179",  # <- A counterfactual non-repeating sequence.
    "01976954310149754605" + "01976954310149754605",  # <- The first example again, but we'll patch this one's activations.
]
```

--------------------------------

### Define a Nested PyTree Object

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/selectors.ipynb

Creates a complex nested Python object containing dictionaries, lists, and JAX arrays, serving as an example for Penzai's selection capabilities.

```python
my_nested_object = {
    "a": 1,
    "b": jnp.arange(10),
    "c": [
        {"value": jnp.arange(12)},
        {"value": jnp.zeros([7])},
        {"value": 3},
    ]
}
```

--------------------------------

### Pretty-print Individual Values with pz.show

Source: https://github.com/google-deepmind/penzai/blob/main/docs/api/treescope.rst

Demonstrates how to pretty-print individual values using the pz.show function. This is a direct way to visualize data structures.

```python
pz.show(value)
```

--------------------------------

### Select and Show Parameters

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Selects and displays all instances of `pz.Parameter` within a given Penzai model. This helps identify learnable parameters.

```python
pz.select(mlp).at_instances_of(pz.Parameter).show_selection()
```

--------------------------------

### Apply Explicit Attention Mask during Model Call

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/howto_reference.md

Example of calling a model modified with `ApplyExplicitAttentionMask`. Ensure inputs like 'tokens', 'token_positions', and 'attn_mask' have the correct named shapes.

```python
# tokens should have named shape {..., "seq": n_seq}
# token_positions should have named shape {..., "seq": n_seq}
# attn_mask should be a boolean array with named shape
#   {..., "seq": n_seq, "kv_seq": n_seq}
token_logits = explicit_attn_model(
  tokens, token_positions=token_positions, attn_mask=attn_mask
)
```

--------------------------------

### JIT-compiling Model Initializer

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/jitting_and_sharding.ipynb

Shows how to JIT-compile a model initializer function using the functional API. This function creates a sequential model with linear and counter layers and returns unbound, frozen variables.

```python
@jax.jit
def functional_init(init_base_rng):
  model = pz.nn.Sequential([
      pz.nn.Linear.from_config(
          name="linear",
          init_base_rng=init_base_rng,
          input_axes={"features": 8},
          output_axes={"features_out": 8},
      ),
      CounterLayer(counter=pz.StateVariable(value=0, label="counter")),
  ])
  # Unbind and also freeze all variables:
  return pz.unbind_variables(model, freeze=True)
```

--------------------------------

### Defining a PositionalSharding

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/jitting_and_sharding.ipynb

Creates a `PositionalSharding` object using a device mesh. This sharding type maps array dimensions directly to device dimensions, useful for basic distributed array setup.

```python
pos_sharding = jax.sharding.PositionalSharding(devices)
```

```python
pos_sharding
```

--------------------------------

### Visualize Learned Parameters

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Displays the learned parameters of the LoRA adapter using Penzai's autovisualizer. This helps confirm that the model has learned meaningful updates, especially by observing non-zero 'B' matrices.

```python
%%autovisualize
_, params = pz.unbind_params(loraified_gemma_model)
params
```

--------------------------------

### Modify Nested Layer Instances

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/howto_reference.md

Chain `pz.select` calls to target specific layer types within nested structures. This example modifies linear layers found within query heads of attention mechanisms.

```python
( \
    pz.select(model) \
    .at_instances_of(pz.nn.Attention) \
    .at(lambda attn: attn.input_to_query) \
    .at_instances_of(pz.nn.Linear) \
    .apply(some_func)
)
```

--------------------------------

### Standard from_config Signature

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

Illustrates the conventional signature for Penzai layer builders like `from_config`. It includes a name for parameter uniqueness and an RNG for initialization.

```python
def from_config(cls, name: str, init_base_rng: jax.Array | None, ...):
  ...
```

--------------------------------

### Rewire Model for Activation Patching

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Applies rewiring to the model's value projections in specified transformer blocks and attention heads. This setup directs activations from different 'worlds' based on induction head masks.

```python
# Start with the original unmodified model checkpoint.
rewired_model = model

# Rewire the attention values in block 20 and 21:
for block_index, induction_head_mask in [
    (20, block_20_induction_heads),
     (21, block_21_induction_heads),
]:
  rewired_model = (
      pz.select(rewired_model)
      .at_instances_of(transformer.model_parts.TransformerBlock)
      .assert_count_is(28)
      .pick_nth_selected(block_index)
      .at_instances_of(pz.nn.Attention)
      .at(lambda attn: attn.input_to_value.sublayers[-1])  # <- the value projections
      .assert_count_is(1)
      .insert_after(RewireComputationPaths(
          worlds_axis="worlds",
          world_ordering=world_ordering,
          taking={
              "original": From("original"),
              "nonrepeating": From("nonrepeating"),
              "patched_induction_values": (
                  # The induction heads read from "nonrepeating".
                  From("nonrepeating", weight=pz.nx.nmap(jnp.where)(induction_head_mask, 1., 0.)),
                  # Everything other than the induction heads take values from "patched_induction_values".
                  From("patched_induction_values", weight=pz.nx.nmap(jnp.where)(induction_head_mask, 0., 1.)),
              ),
          },
      ))
  )
```

--------------------------------

### Download Gemma Weights

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Downloads the Gemma 7B model weights from Kaggle and constructs the checkpoint and vocabulary paths.

```python
weights_dir = kagglehub.model_download('google/gemma/Flax/7b')
ckpt_path = os.path.join(weights_dir, '7b')
vocab_path = os.path.join(weights_dir, 'tokenizer.model')
```

--------------------------------

### Initialize Model Parameters with Sharding Inference

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/jitting_and_sharding.ipynb

Uses `sharding_util.sharded_init` to trace a model initializer, infer sharding specifications based on parameter axis names, and then run the initialization.

```python
from penzai.toolshed import sharding_util
from penzai.models.transformer.variants import llamalike_common
```

```python
# Very small transformer config, for demo purposes
config = llamalike_common.LlamalikeTransformerConfig(
    num_kv_heads=2,
    query_head_multiplier=1,
    embedding_dim=64,
    projection_dim=16,
    mlp_hidden_dim=128,
    num_decoder_blocks=2,
    vocab_size=100,
    mlp_variant="geglu_approx",
    rope_wavelength=10_000,
    tie_embedder_and_logits=True,
    use_layer_stack=False,
    parameter_dtype=jnp.float32,
    activation_dtype=jnp.float32,
)

tiny_transformer = sharding_util.sharded_init(
    llamalike_common.build_llamalike_transformer,
    config=config,
    init_base_rng=jax.random.key(42),
    mesh=jax.sharding.Mesh(devices, axis_names=('devices',)),
    axis_name_to_mesh_name={
        # Shard the embedding dimension across devices.
        "embedding": "devices",
    },
)

tiny_transformer
```

--------------------------------

### Analyze Gradients for Induction Heads

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Prints the maximum gradient value for heads not included in the `top_heads_mask` and renders the gradients, highlighting areas where the mask is not applied. This helps identify the most impactful heads when starting from an un-ablated model.

```python
print(jnp.max((accuracy_grads * (1-top_heads_mask)).untag("blocks", "heads").unwrap()))
treescope.render_array(
    accuracy_grads,
    valid_mask=1-top_heads_mask
)
```

--------------------------------

### Calculate Gradients from Ablated State

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Calculates gradients starting from an ablated model state, using `top_heads_mask` to define the ablated condition. This is used to assess the impact of adding back specific heads to an already ablated model.

```python
accuracy_grads = pz.variable_jit(jax.grad(get_ablated_avg_log_prob, argnums=0))(
    top_heads_mask,
    model,
    token_seq,
    loss_mask=(pz.nx.arange("seq", 40) > 21)
)
accuracy_grads
```

--------------------------------

### Calling JIT-compiled Initializer

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/jitting_and_sharding.ipynb

This snippet demonstrates calling the `functional_init` function to obtain a model with slots and initial variable values.

```python
model_with_slots, init_var_values = functional_init(jax.random.PRNGKey(0))
```

--------------------------------

### Calculate Gradients for Ablated Model

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Calculates the gradients of the log probability with respect to the model's parameters, focusing on improvements to the second repetition of a sequence. This is used to analyze the impact of specific heads when starting from an ablated model state.

```python
accuracy_grads = pz.variable_jit(jax.grad(get_ablated_avg_log_prob, argnums=0))(
    pz.nx.ones({"blocks": 28, "heads": 16}),
    model,
    token_seq,
    # Focus on improvements to the second repetition
    loss_mask=(pz.nx.arange("seq", 40) > 21)
)
accuracy_grads
```

--------------------------------

### Initialize MLP Model from Scratch

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/howto_reference.md

Initialize an MLP model using its constructor classmethod `from_config`. Pass a name and a JAX PRNG key for initialization. Set `init_base_rng=None` to build the model without initializing parameters.

```python
mlp = simple_mlp.MLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(10),
    feature_sizes=[2, 32, 32, 2],
)
```

--------------------------------

### Visualize RewireComputationPaths Instances

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Selects and visualizes all instances of `RewireComputationPaths` within a given model. This is useful for verifying that the rewiring logic has been applied correctly.

```python
pz.select(rewired_model).at_instances_of(RewireComputationPaths)
```

--------------------------------

### Configure Sharding for Model Restoration

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Sets up device mesh and sharding configurations for restoring model parameters across multiple JAX devices.

```python
n_devices = jax.local_device_count()
sharding_devices = mesh_utils.create_device_mesh((n_devices,))
sharding = jax.sharding.PositionalSharding(sharding_devices)
restore_args = jax.tree_util.tree_map(
    lambda m: orbax.checkpoint.ArrayRestoreArgs(
        restore_type=jax.Array,
        sharding=sharding.reshape((1,) * (len(m.shape) - 1) + (n_devices,)))
    ,
    metadata,
)
flat_params = checkpointer.restore(ckpt_path, restore_args=restore_args)
```

--------------------------------

### Untagging and Applying Softmax with Penzai

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

This example shows how to remove a name from an axis, apply a JAX function (like softmax) using `pz.nx.nmap` which vectorizes over named axes, and then re-assign a name to the resulting axis. This is a common pattern for applying operations along specific named dimensions.

```python
# Un-tag the vocabulary axis:
untagged = named.untag("vocabulary")
# Map the ordinary JAX softmax function over the temporary positional axis:
softmaxed = pz.nx.nmap(jax.nn.softmax)(untagged, axis=0)
# Tag the positional axis with a name again:
softmaxed.tag("vocabulary")
```

--------------------------------

### Swapping Attention Layers with KVCachingAttention

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

Demonstrates how to use `pz.select` to replace `Attention` layers with `KVCachingAttention` layers, enabling KV caching.

```python
(pz.select(model)
 .at_instances_of(pz.nn.Attention)
 .apply(lambda attn: pz.nn.KVCachingAttention.from_uncached(attn, **kwargs)))
```

--------------------------------

### Configure and Apply Rewiring for Path Analysis

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Sets up a `RewireComputationPaths` object to define how different worlds read from each other, then integrates this rewirer into the model. Finally, it applies the ablation mask to the rewired model.

```python
read_rewirer = RewireComputationPaths(
    worlds_axis="worlds",
    world_ordering=world_ordering,
    taking={
        "original": From("original"),
        "original_ablate_direct": From("original"),
        "fully_ablated": From("fully_ablated"),
        "ablated_restore_direct": From("fully_ablated"),
    },
)
rewired_model = (
    pz.select(model)
    .at_instances_of(pz.nn.Residual)
    .at(lambda r: r.delta.sublayers[0])  # <- assuming each residual contains a Sequential
    .insert_before(read_rewirer)
)
ablated_rewired_model = knock_out_heads(rewired_model, per_world_head_mask)

read_rewirer.path_matrix()
```

--------------------------------

### Import Penzai and Model Components

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

Imports Treescope for visualization and specific components from the Penzai library, including the pz module and simple_mlp models.

```python
import treescope
import penzai
from penzai import pz
from penzai.models import simple_mlp
```

--------------------------------

### Instantiating and calling a custom layer

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Instantiate the `DisplayIntermediateValue` layer and call it with a sample input, then display the final output.

```python
layer = DisplayIntermediateValue()
output = layer(123)
pz.show("Final:", output)
```

--------------------------------

### Define Custom Layer by Subclassing Sequential

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/howto_reference.md

Create a reusable or semantically meaningful model component by subclassing `pz.nn.Sequential`. Define a custom `from_config` class method for initialization, following the convention used by Penzai's built-in layers.

```python
@pz.pytree_dataclass(has_implicitly_inherited_fields=True)
class MyCustomLayer(pz.nn.Sequential):
  @classmethod
  def from_config(cls, name: str, init_base_rng: jax.Array | None, ...):
    ...
```

--------------------------------

### Initialize SimpleMLP Model with Parameters

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

Initializes the parameters of a SimpleMLP model as mutable Variable objects using a provided PRNGKey. This prepares the model for training.

```python
model = SimpleMLP.from_config(
    name="mlp",
    init_base_rng=jax.random.key(42),
    feature_sizes=[8, 32, 32, 8],
    activation=jax.nn.relu,
    features_axis="features",
)
model
```

--------------------------------

### Import Standard Python Libraries

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Imports essential Python libraries for numerical computation, machine learning, and checkpointing. These are standard imports for Penzai-based projects.

```python
from __future__ import annotations

import os
import gc

import jax
import jax.numpy as jnp
import numpy as np
import orbax.checkpoint
import optax
from jax.experimental import mesh_utils
```

--------------------------------

### Load GPT-NeoX / Pythia Model from HuggingFace

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/howto_reference.md

Loads a GPT-NeoX or Pythia model from HuggingFace transformers and converts it to a Penzai model.

```python
import transformers
from penzai.models.transformer import variants

# To load a GPT-NeoX / Pythia model:
hf_model = transformers.GPTNeoXForCausalLM.from_pretrained(...)
pz_model = variants.gpt_neox.gpt_neox_from_huggingface_model(hf_model)
```

--------------------------------

### Create Selection from a Struct

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/selectors.ipynb

Shows how to create a selection using the `.select()` method, which is available on objects that are subclasses of `pz.Struct`.

```python
@pz.pytree_dataclass
class MyStruct(pz.Struct):
  foo: typing.Any

MyStruct(4).select()
```

--------------------------------

### Loading Pretrained Gemma Transformer (V1 vs V2)

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/v2_differences.md

Migrate from the V1 method of loading Gemma models to the V2 API, which uses a new transformer implementation and more generic component types.

```python
# Old
from penzai.deprecated.v1.example_models import gemma
model = gemma.model_core.GemmaTransformer.from_pretrained(flax_params_dict)
# (model is an instance of GemmaTransformer)

# New
from penzai.models.transformer import variants
model = variants.gemma.gemma_from_pretrained_checkpoint(flax_params_dict)
# (model is an instance of TransformerLM)
```

--------------------------------

### Configure JAX to Use Multiple CPU Devices

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/jitting_and_sharding.ipynb

Forces the JAX CPU backend to simulate multiple devices, useful for testing distributed or multi-device functionality on a single machine. This should be set before JAX initializes devices.

```python
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"
```

--------------------------------

### Import Penzai Model and Toolshed Modules

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Imports specific modules from Penzai for transformer and simple MLP models, as well as tools for token visualization, basic training, and JIT wrapping.

```python
from penzai.models import transformer
from penzai.models import simple_mlp
from penzai.toolshed import token_visualization
from penzai.toolshed import basic_training
from penzai.toolshed import jit_wrapper
```

--------------------------------

### MLP Layer Configuration in Penzai v2

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/v2_differences.md

Shows how the `from_config` method for an MLP layer is updated in v2 to accept `name` and `init_base_rng` for eager parameter initialization.

```diff
 @pz.pytree_dataclass(has_implicitly_inherited_fields=True)
 class MLP(pz.nn.Sequential):
   """Sequence of Affine layers."""

   @classmethod
   def from_config(
       cls,
+      name: str,
+      init_base_rng: jax.Array | None,
       feature_sizes: list[int],
       activation_fn: Callable[[jax.Array], jax.Array] = jax.nn.relu,
       feature_axis: str = "features",
   ) -> MLP:
     assert len(feature_sizes) >= 2
     children = []
     for i, (feats_in, feats_out) in enumerate(
         zip(feature_sizes[:-1], feature_sizes[1:])
     ):
       if i:
         children.append(pz.nn.Elementwise(activation_fn))
       children.append(
-          pz.nn.add_parameter_prefix(
-              f"Affine_{i}",
-              pz.nn.Affine.from_config(
-                  input_axes={feature_axis: feats_in},
-                  output_axes={feature_axis: feats_out},
-              ),
+          pz.nn.Affine.from_config(
+              name=f"{name}/Affine_{i}",
+              init_base_rng=init_base_rng,
+              input_axes={feature_axis: feats_in},
+              output_axes={feature_axis: feats_out},
            )
        )
     return cls(sublayers=children)
```

--------------------------------

### Import Core Libraries

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Imports essential libraries for numerical computation, model checkpointing, and distributed array handling.

```python
from __future__ import annotations
from typing import Any

import os
import dataclasses
import gc

import jax
import jax.numpy as jnp
import numpy as np
import orbax.checkpoint
from jax.experimental import mesh_utils
```

--------------------------------

### Constructing a Simple MLP Model

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

This code shows how to construct a simple Multi-Layer Perceptron (MLP) model in Penzai, defining its layers and their parameters.

```python
copied = penzai.models.simple_mlp.MLP( # Sequential
  sublayers=[
    penzai.nn.linear_and_affine.Affine( # Sequential
      sublayers=[
        penzai.nn.linear_and_affine.Linear(weights=penzai.core.variables.ParameterSlot(label='mlp/Affine_0/Linear.weights'), in_axis_names=('features',), out_axis_names=('features_out',)),
        penzai.nn.linear_and_affine.RenameAxes(old=('features_out',), new=('features',)),
        penzai.nn.linear_and_affine.AddBias(bias=penzai.core.variables.ParameterSlot(label='mlp/Affine_0/AddBias.bias'), new_axis_names=()),
      ],
    ),
    penzai.nn.basic_ops.Elementwise(fn=jax.nn.relu),
    penzai.nn.linear_and_affine.Affine( # Sequential
      sublayers=[penzai.nn.linear_and_affine.Linear(weights=penzai.core.variables.ParameterSlot(label='mlp/Affine_1/Linear.weights'), in_axis_names=('features',), out_axis_names=('features_out',)), penzai.nn.linear_and_affine.RenameAxes(old=('features_out',), new=('features',)), penzai.nn.linear_and_affine.AddBias(bias=penzai.core.variables.ParameterSlot(label='mlp/Affine_1/AddBias.bias'), new_axis_names=())],
    ),
    penzai.nn.basic_ops.Elementwise(fn=jax.nn.relu),
    penzai.nn.linear_and_affine.Affine( # Sequential
      sublayers=[penzai.nn.linear_and_affine.Linear(weights=penzai.core.variables.ParameterSlot(label='mlp/Affine_2/Linear.weights'), in_axis_names=('features',), out_axis_names=('features_out',)), penzai.nn.linear_and_affine.RenameAxes(old=('features_out',), new=('features',)), penzai.nn.linear_and_affine.AddBias(bias=penzai.core.variables.ParameterSlot(label='mlp/Affine_2/AddBias.bias'), new_axis_names=())],
    ),
  ],
)

copied == unbound_mlp
```

--------------------------------

### Import Statements for Penzai V1 and V2 APIs

Source: https://github.com/google-deepmind/penzai/blob/main/docs/guides/v2_differences.md

Illustrates how to import components from both the legacy V1 API and the new V2 API in Penzai. Use the `penzai.deprecated.v1` namespace for V1 components.

```python
# Old V1 API:
from penzai.deprecated.v1 import pz
from penzai.deprecated.v1.example_models import simple_mlp
import penzai.deprecated.v1.toolshed

# New V2 API:
from penzai import pz
from penzai.models import simple_mlp
import penzai.toolshed
```

--------------------------------

### Initialize Orbax Checkpointer

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/lora_from_scratch.ipynb

Initializes the Orbax checkpointer to manage model checkpoints.

```python
checkpointer = orbax.checkpoint.PyTreeCheckpointer()
metadata = checkpointer.metadata(ckpt_path)
```

--------------------------------

### Load Gemma Model Checkpoint

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Builds a Gemma model from a pre-trained checkpoint. Activations are configured to be computed in float32 precision, while weights remain in bfloat16.

```python
model = transformer.variants.gemma.gemma_from_pretrained_checkpoint(
    flat_params,
    upcast_activations_to_float32=True
)
```

--------------------------------

### Define a SimpleMLP Model Configuration

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

Builds the structure of a SimpleMLP model without initializing parameters. Use this to define the network architecture.

```python
SimpleMLP.from_config(
    name="mlp",
    init_base_rng=None,
    feature_sizes=[8, 32, 32, 8],
    activation=jax.nn.relu,
    features_axis="features",
)
```

--------------------------------

### Using pz.variable_jit for JIT compilation

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/jitting_and_sharding.ipynb

A wrapper around `jax.jit` that allows function arguments to include `pz.Parameter` and `pz.StateVariable`. It handles updating their values automatically. Note that it does not support returning variables or closing over global references.

```python
@pz.variable_jit
def jitted_call(model, arg):
  return model(arg)
```

```python
jitted_call(model, pz.nx.ones({"features": 8}))
```

```python
jitted_call(model, pz.nx.ones({"features": 8}))
```

--------------------------------

### Visualize Digit Token Probabilities

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Calculates and visualizes the probability distribution of digit tokens from model logits. Requires Penzai and JAX libraries.

```python
all_probs = pz.nx.nmap(jax.nn.softmax)(paired_logits.untag("vocabulary")).tag("vocabulary")
digit_token_ids = pz.nx.wrap(vocab.EncodeAsIds("0123456789")).tag("digits")
digit_probs = all_probs[{"vocabulary": digit_token_ids}]

treescope.render_array(digit_probs, axis_item_labels={"worlds": world_ordering}, vmax=1)
```

--------------------------------

### Show Token Array Visualization

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Uses Penzai's token_visualization utility to display a named array of tokens along with the vocabulary for context.

```python
from penzai.toolshed import token_visualization
token_visualization.show_token_array(token_seq, vocab)
```

--------------------------------

### Visualize Rewired Computation Paths

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/induction_heads.ipynb

Selects instances of RewireComputationPaths within the rewired model for visualization. This helps verify that the rewiring has been applied correctly.

```python
%%autovisualize None
pz.select(rewired_model).at_instances_of(RewireComputationPaths)
```

--------------------------------

### Running a Patched Model

Source: https://github.com/google-deepmind/penzai/blob/main/notebooks/how_to_think_in_penzai.ipynb

Executes the modified Penzai model (`patched`) which now includes the custom `HelloWorld` layer. This demonstrates the effect of the insertion.

```python
# Run it:
patched(pz.nx.ones({"features": 8}))
```