AddBias

class penzai.nn.linear_and_affine.AddBias

Bases: Layer

Shifts an input by a learnable offset (a bias).

This layer uses named arrays to automatically apply across the correct set of dimensions.

Variables:

• bias (parameters.ParameterLike[NamedArray]) – The learnable bias.

• new_axis_names (tuple[str, ...]) – The new axes in the output that we do not expect to see in the input.

Methods

__init__(bias, new_axis_names)
from_config(biased_axes[, new_output_axes, ...])	Constructs an AddBias layer from a configuration.
input_structure()
output_structure()
treescope_color()
__call__(value)	Adds a learned bias to the value.

Attributes

bias
new_axis_names

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.
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.
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: NamedArray) → NamedArray

Adds a learned bias to the value.

classmethod from_config(biased_axes: dict[str, int], new_output_axes: dict[str, int] | None = None, initializer: BiasInitializer = <function zeros>, dtype: jax.typing.DTypeLike = <class 'jax.numpy.float32'>) → AddBias

Constructs an AddBias layer from a configuration.

Parameters:

• biased_axes – Names and lengths for the axes in the input that the bias should act over. Other axes will be broadcast over.

• new_output_axes – Names and lengths of new axes that should be introduced into the input.

• initializer – Function to use to initialize the weight.

• dtype – Dtype for the bias.

Returns:

A new AddBias layer with an uninitialized bias parameter.