AddBias

AddBias#

class penzai.nn.linear_and_affine.AddBias[source]#

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:

Methods

__init__(bias, new_axis_names)

from_config(name, init_base_rng, biased_axes)

Constructs an AddBias layer from a configuration.

treescope_color()

__call__(value, **_unused_side_inputs)

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.

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: NamedArray, **_unused_side_inputs) → NamedArray[source]#

Adds a learned bias to the value.

classmethod from_config(name: str, init_base_rng: jax.Array | None, biased_axes: dict[str, int], new_output_axes: dict[str, int] | None = None, initializer: LinearOperatorWeightInitializer = <function constant_initializer.<locals>._initializer>, dtype: jax.typing.DTypeLike = <class 'jax.numpy.float32'>) → AddBias[source]#

Constructs an AddBias layer from a configuration.

Parameters:

  • name – The name of the layer.

  • init_base_rng – The base RNG to use for initializing model 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. Only the output axes will be set.

  • dtype – Dtype for the bias.

Returns:

A new AddBias layer with an uninitialized bias parameter.

Contents