def full(
    named_shape: Mapping[AxisName, int],
    fill_value: jax.typing.ArrayLike,
    dtype: np.DTypeLike | None = None,
) -> NamedArray:
  """Constructs a full named array with a given shape.

  Args:
    named_shape: Named shape for the result.
    fill_value: Value to fill the array with.
    dtype: Optional dtype for the result. If not provided, will be inferred from
      ``fill_value``.

  Returns:
    NamedArray with the given named shape, filled with ``fill_value``.
  """
  return NamedArray(
      named_axes=collections.OrderedDict(named_shape),
      data_array=jnp.full(tuple(named_shape.values()), fill_value, dtype=dtype),
  )
Constructs a full named array with a given shape.

  Args:
    named_shape: Named shape for the result.
    fill_value: Value to fill the array with.
    dtype: Optional dtype for the result. If not provided, will be inferred from
      ``fill_value``.

  Returns:
    NamedArray with the given named shape, filled with ``fill_value``.
  
penzai/core/named_axes.py