def wrap(cls, array: jax.typing.ArrayLike, *names: AxisName) -> NamedArray:
    """Wraps a positional array as a ``NamedArray``.

    Args:
      array: Array to wrap.
      *names: Optional names for the axes of the array. If provided, must be the
        same length as the array's shape. This is a convenience wrapper so that
        you can call ``wrap(array, "foo", "bar")`` instead of
        ``wrap(array).tag("foo", "bar")``.

    Returns:
      An equivalent ``NamedArray`` for the given array. If ``names`` is
      provided, the resulting array will have those names assigned to the
      corresponding axes. Otherwise, the resulting array will have a positional
      shape.
    """
    wrapped = NamedArray(
        named_axes=collections.OrderedDict(), data_array=jnp.asarray(array)
    )
    if names:
      return wrapped.tag(*names)
    else:
      return wrapped
Wraps a positional array as a ``NamedArray``.

    Args:
      array: Array to wrap.
      *names: Optional names for the axes of the array. If provided, must be the
        same length as the array's shape. This is a convenience wrapper so that
        you can call ``wrap(array, "foo", "bar")`` instead of
        ``wrap(array).tag("foo", "bar")``.

    Returns:
      An equivalent ``NamedArray`` for the given array. If ``names`` is
      provided, the resulting array will have those names assigned to the
      corresponding axes. Otherwise, the resulting array will have a positional
      shape.
    
penzai/core/named_axes.py