DOCSEARCHbyAlgolia
Submit Docs

Penzai

repository·main·Indexed 23 days ago

https://github.com/google-deepmind/penzai

A JAX research toolkit for building, editing, and visualizing neural networks. Penzai treats models as legible, functional pytree data structures to facilitate model interpretability, surgery, and analysis. Key features include the pz.select utility for type-driven pytree traversals, pz.StateVariable for capturing intermediate activations, and pz.nx for named axis vectorization. The library includes a comprehensive set of neural network building blocks in pz.nn, such as Linear, LayerNorm, and various combinators like Sequential and Residual.

Tokens
30.1K
Snippets
74
Records
142
Agent score
83%

What's inside penzai

  1. Use pz.nn for neural network construction

    main
    The penzai.pz.nn namespace provides a collection of building blocks for defining neural networks. It includes standard layers (Linear, LayerNorm), combinators (Sequential, Residual), and specialized modules for language modeling (Attention, Embedding).

  2. Understand the differences between Penzai V1 and V2 APIs

    main

    Penzai has transitioned from a V1 API to a V2 API. This change is a breaking change introduced in version v0.2.0.

    API Locations

    • V2 API (Current): Available in penzai.nn and used in penzai.models. This is the simplified, recommended design.
    • V1 API (Deprecated): Located in penzai.deprecated.v1.nn and penzai.deprecated.v1.data_effects. Use this only if maintaining legacy code.

    Key Conceptual Shifts

    FeatureV1 API ApproachV2 API Approach
    Parameters & StateImmutable PyTrees of arrays; required data_effects and handler blocks for sharing/state.Mutable Parameter and StateVariable objects using standard Python shared-reference semantics.
    Side InputsInjected as attributes via effect handlers.Passed as keyword arguments (**kwargs) to a layer's __call__ method.
    InitializationRequired a separate pz.nn.initialize_parameters step.Eager initialization; models are initialized directly.
    Parameter SharingRequired explicit 'lookup' effects and handler invariants.Works by default; shared parameters are multiple references to the same Parameter object.
    Data EffectsUsed penzai.deprecated.v1.data_effects to manage side effects.Deprecated. Use Parameter, StateVariable, or keyword arguments instead.

  3. Use pz.select to modify model architectures

    main

    The pz.select utility (the pz.select module) acts as a pytree 'swiss-army-knife'. It generalizes JAX's .at[...].set(...) syntax to arbitrary type-driven traversals. This allows you to perform complex rewrites, such as inserting new layers after specific existing layers.

    # Example: Inserting a layer after all Elementwise nonlinearities
saving_model = (
    pz.select(mlp)
    .at_instances_of(pz.nn.Elementwise)
    .insert_after(AppendIntermediate(var))
)

  4. Manage Parameters and State Variables

    main

    Penzai handles mutable state by embedding stateful parameters and variables into JAX PyTrees. This allows state to be passed across JAX transformation boundaries (like jit).

    Key Components

    • pz.Parameter / pz.ParameterValue / pz.ParameterSlot: For handling model weights and parameters.
    • pz.StateVariable / pz.StateVariableValue / pz.StateVariableSlot: For handling mutable state (e.g., running averages).
    • pz.RandomStream: For managing stochasticity.

    State Manipulation Utilities

    Use these to manage how state is handled during JAX transformations:

    • pz.bind_variables / pz.unbind_variables: Bind or unbind state variables.
    • pz.freeze_variables / pz.unbind_variables: Freeze or unbind state.
    • pz.bind_params / pz.unbind_params: Bind or unbind parameters.
    • pz.freeze_params / pz.unbind_params: Freeze or unbind parameters.
    • pz.variable_jit: A utility to assist with JIT-compiling functions containing variables.

  5. Understand the difference between Parameter and StateVariable in V2

    main

    In the Penzai V2 API, parameters and state variables are no longer ordinary PyTree nodes but are represented as mutable "variable" objects embedded within the model structure. This allows for sharing values between multiple parts of a model using standard Python reference semantics.

    There are two primary types of variables:

    1. Parameter: Used for model parameters that are updated by an optimizer but are not modified during the forward pass. They can be shared between multiple models or model components.
    2. StateVariable: A mutable variable intended to be modified during the forward pass.

    Important Note on JAX Compatibility: Because these variable objects are mutable, they cannot be passed directly through most JAX transformations. You must use Penzai utilities to identify/extract variables and run model logic in a functional way. To support this, every variable object must have a unique label (either manually specified or automatically generated).

    While the variables themselves are mutable leaves, the model layers remain immutable PyTree nodes, allowing for safe model copying and modification.

  6. Vectorize operations over named axes with pz.nx.nmap

    main

    To perform operations on specific axes of named arrays, use pz.nx.nmap. This functions similarly to jax.vmap but uses axis names for inference instead of positional indices.

    Behavior:

    • Every axis in the .named_shape of the input arguments is vectorized over. Axes with matching names are paired.
    • Every axis in the .positional_shape is preserved and visible inside the mapped function.

    Pattern for positional operations:

    1. .untag() the desired axes to move them to .positional_shape.
    2. Call pz.nx.nmap(operation)(...args...).
    3. .tag() the resulting axes back to the .named_shape.

  7. How parameters and state work in the V2 API

    main

    In the V2 API, parameters and state are no longer just JAX arrays within an immutable PyTree. Instead, they are represented by mutable Parameter and StateVariable objects.

    • Model Structure: Model layers remain immutable JAX PyTree nodes. However, the leaves of these trees can now be Parameter or StateVariable instances instead of raw JAX arrays.
    • Sharing Semantics: Because these are Python objects, they follow ordinary Python shared-reference semantics. If multiple layers hold a reference to the same Parameter object, they share that parameter automatically.
    • Functional Usage: While the objects themselves are mutable, Penzai provides helper functions to manipulate these variables and call models purely functionally when needed.
    • Benefit: This design removes the need for the complex 'effect handler' boilerplate required in V1 to manage state and parameter sharing.

  8. How V1 Layers and Parameters work

    main

    In the V1 API, most models and layers are subclasses of pz.Layer, which can be called with a single argument.

    To enable runtime shape-checking on pz.Layer subclasses, use the @pz.checked_layer_call or @pz.unchecked_layer_call decorators.

    Parameters in V1 are ordinary PyTree nodes. Parameter sharing is managed via metadata and PyTree transformations within the model using utilities like pz.nn.mark_shareable and pz.nn.attach_shared_parameters.

  9. Build neural networks with pz.nn

    main
    The pz.nn namespace is a declarative neural network system. It uses a combinator-based design where all model operations are exposed as nodes in the model's PyTree. This makes the entire architecture inspectable and manipulatable as a data structure. pz.nn re-exports layers from penzai.nn into a single namespace.

  10. Build models using combinators and primitives

    main

    Penzai follows a compositional design where complex architectures are built by combining combinators (which define logic flow) and primitives (which define basic operations).

    Combinators:

    • pz.nn.Sequential: Runs layers in sequence.
    • pz.nn.Residual: Runs child layers and adds their output to the input.
    • pz.nn.BranchAndAddTogether / pz.nn.BranchAndMultiplyTogether: Combines outputs of different components.
    • pz.nn.Attention: Routes inputs between query, key, value, masking, and output.

    Primitives:

    • pz.nn.Affine, pz.nn.Linear, pz.nn.AddBias, pz.nn.Elementwise, pz.nn.Softmax, etc.

  11. How to pass side inputs in the V2 API

    main

    In the V2 API, side inputs (data that is not the primary argument but is needed by a layer) should be passed as keyword arguments to the __call__ method of each layer.

    Implementation Details

    • Signature Change: The signature of Layer.__call__ has changed from __call__(self, arg, /) to __call__(self, arg, /, **kwargs).
    • Layer Responsibility: Layers are expected to ignore any keyword arguments they do not explicitly recognize. This allows you to pass a common set of side inputs through a deep model without every intermediate layer needing to implement logic for them.

  12. Inspect NamedArray shapes

    main

    Named arrays maintain two distinct shape representations:

    • .positional_shape: A sequence of dimension sizes for axes that are not named.
    • .named_shape: A dictionary mapping axis names to their dimension sizes.

    Note that an axis exists in either the positional shape or the named shape, but never both. Calling .untag() moves axes from the named shape to the positional shape.