def from_config(
      cls,
      name: str,
      init_base_rng: jax.Array | None,
      biased_axes: dict[str, int],
      new_output_axes: dict[str, int] | None = None,
      initializer: LinearOperatorWeightInitializer = zero_initializer,
      dtype: jax.typing.DTypeLike = jnp.float32,
  ) -> AddBias:
    """Constructs an ``AddBias`` layer from a configuration.

    Args:
      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.
    """
    if new_output_axes is None:
      new_output_axes = {}

    return cls(
        bias=parameters.make_parameter(
            f"{name}.bias",
            init_base_rng,
            initializer,
            input_axes={},
            output_axes={**biased_axes, **new_output_axes},
            parallel_axes={},
            convolution_spatial_axes={},
            dtype=dtype,
        ),
        new_axis_names=tuple(new_output_axes.keys()),
    )
Constructs an ``AddBias`` layer from a configuration.

    Args:
      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.
    
penzai/nn/linear_and_affine.py