def constant_initializer(value: float) -> LinearOperatorWeightInitializer:
  """Returns an initializer that uses a constant value."""

  def _initializer(
      key: jax.Array,
      *,
      input_axes: dict[str, int],
      output_axes: dict[str, int],
      parallel_axes: dict[str, int],
      convolution_spatial_axes: dict[str, int],
      dtype: jax.typing.DTypeLike,
  ) -> jax.Array:
    """Zeros initializer for named arrays."""
    del key
    check_unique_axis_names_for_initializer(
        input_axes, output_axes, parallel_axes, convolution_spatial_axes
    )
    return named_axes.full(
        {
            **input_axes,
            **output_axes,
            **parallel_axes,
            **convolution_spatial_axes,
        },
        value,
        dtype=dtype,
    )

  return _initializer
Returns an initializer that uses a constant value.
penzai/deprecated/v1/nn/linear_and_affine.py