Checks if the type of a node is exactly `Sequential`, not a subclass.

Checks if the type of a node is exactly `Sequential`, not a subclass.

Returns True if there are no variables.

Returns whether we are currently rendering a selection.

Returns a custom GetAttrKey with layer stack metadata.

Extracts the stacked axes from a keypath. This can be used to initialize new variables for transformations that modify layers inside a LayerStack

Resolve a GitHub URL corresponding to Python object.

Converts a HuggingFace Llama model to a Penzai model. This function converts Llama models from their HuggingFace implementations to Penzai. (Othe

Converts a "llama-like" HuggingFace model to a Penzai model. This function converts Llama-like models from their HuggingFace implementations to P

Replaces Linear layers inside a selected part of a model with LoRA blocks. This function should usually be called after freezing the existing weigh

Replaces Linear layers inside a selected part of a model with LoRA blocks. This function should usually be called after freezing the existing weigh

Makes a parameter variable (or slot) with a given name and initializer. Args: name: The name of the parameter. init_base_rng: The base PRNG

Marks all uninitialized parameters in ``submodel`` as shareable. This function is used to annotate the parameters in part of a model as being sha

Constructs either a stochastic or disabled dropout layer. Args: drop_rate: Probability of dropping an element. If None, dropout will be d

Constructs either a stochastic or disabled dropout layer. Args: drop_rate: Probability of dropping an element. If None, dropout will be d

Converts a HuggingFace Mistral model to a Penzai model. This function converts Mistral models from their HuggingFace implementations to Penzai. (

The full model, including parameters and nonlearnable parts.

Shards a tree of `pz.nx.NamedArray` objects based on their axis names. Args: tree: A PyTree of NamedArrays. mesh: The Mesh to shard the tre

A mapping of axis names to their sizes.

Constructs a named array of ones with a given shape. Args: named_shape: Named shape for the result. dtype: Optional dtype for the result.

The axis names and sizes that will appear in the output only.

The axis names and sizes that will appear in the output only.

The axis names and sizes that should appear in both input and output.

The axis names and sizes that should appear in both input and output.

Patches the IPython debugger to use `IPython.display.display`.

A tuple of axis sizes for any anonymous axes.

Constructs a pretty name from a keypath and an object. This can be used to construct human-readable names for locations inside a PyTree. Args:

Processes leaves or values where filter_fn is True.

Process PyTree leaves or subtrees where _is_hole_or_quote is True.

Processes one of the selected subtrees.

Recursively walks subtrees one level at a time. Args: keypath: Keypath to this subtree leaf_or_subtree: The subtree

Splits a PRNG key into a `NamedArray` of PRNG keys with the given names. Args: key: PRNG key to split. Can also be a NamedArray of keys with di