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.
bind_variables(variables[, allow_unused])	Convenience function to bind variables to a layer.
from_attributes(**field_values)	Directly instantiates a struct given all of its fields.
key_for_field(field_name)	Generates a JAX PyTree key for a given field name.
select()	Wraps this struct in a selection, enabling functional-style mutations.
stateless_call(variable_values, argument, /, ...)	Calls a layer with temporary variables, without modifying its state.
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, **side_inputs)	Runs each of the sublayers in sequence.