SupportsParameterRenaming

class penzai.nn.parameters.SupportsParameterRenaming

Bases: ABC

Base class that identifies a PyTree node as supporting parameter renaming.

This class can be used to identify nodes that have parameter names, and which should be renamed when add_parameter_prefix is called. Subclassing this allows new parameter variants to also support renaming.

Methods

with_renamed_parameters(rename_fn)

Returns a copy of this node with parameterss renamed.

Inherited Methods

(expand to view inherited methods)

__init__()

abstract with_renamed_parameters(rename_fn: Callable[[str], str])

Returns a copy of this node with parameterss renamed.

If implemented on a parameter, this should make a copy of the parameter with a new name. If implemented on a container of parameters, this should rename every parameter, possibly by recursively calling with_renamed_parameters on its children.

Parameters:

rename_fn – Function mapping old names to new names.

Returns:

Copy of self with renaming applied.