class penzai.nn.parameters.SupportsParameterRenaming[source]#

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.



Returns a copy of this node with parameterss renamed.

Inherited Methods

(expand to view inherited methods)


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

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.


rename_fn – Function mapping old names to new names.


Copy of self with renaming applied.