LinearInPlace

class penzai.nn.linear_and_affine.LinearInPlace

Bases: Sequential

Container for “in-place” linear operators that preserve axis names.

This is used when initializing Linear layers that have overlapping names in their input and output specifications. We subclass Sequential to make this layer type easier to identify and manipulate.

Methods

__init__(sublayers)
treescope_color()

Attributes

sublayers

Inherited Methods

(expand to view inherited methods)

attributes_dict()	Constructs a dictionary with all of the fields in the class.
from_attributes(**field_values)	Directly instantiates a struct given all of its fields.
input_structure()	Returns the input structure of this layer.
key_for_field(field_name)	Generates a JAX PyTree key for a given field name.
output_structure()	Returns the output structure of this layer.
select()	Wraps this struct in a selection, enabling functional-style mutations.
tree_flatten()	Flattens this tree node.
tree_flatten_with_keys()	Flattens this tree node with keys.
tree_unflatten(aux_data, children)	Unflattens this tree node.
__call__(value)	Runs each of the sublayers in sequence.