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

  Args:
    named_shape: Named shape for the result.
    dtype: Optional dtype for the result.

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

  Args:
    named_shape: Named shape for the result.
    dtype: Optional dtype for the result.

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