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

    Args:
      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.
    """
    if new_output_axes is None:
      new_output_axes = {}
    bias_shape = []
    bias_names = []
    for name, size in itertools.chain(
        biased_axes.items(), new_output_axes.items()
    ):
      if name in bias_names:
        raise ValueError(f"Duplicate axis name in bias: {name}")
      bias_names.append(name)
      bias_shape.append(size)

    def bias_initializer(prng_key):
      return named_axes.wrap(
          initializer(prng_key, tuple(bias_shape), dtype)
      ).tag(*bias_names)

    return cls(  # pytype: disable=wrong-arg-types
        bias=parameters.UninitializedParameter(bias_initializer, name="bias"),
        new_axis_names=tuple(new_output_axes.keys()),
    )
Constructs an ``AddBias`` layer from a configuration.

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