Affine#
- class penzai.nn.linear_and_affine.Affine[source]#
Bases:
Sequential
Affine layer: combination of
Linear
andAddBias
.Methods
__init__
(sublayers)from_config
(name, init_base_rng, input_axes, ...)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.