wrap

penzai.core.named_axes.wrap(array: jax.typing.ArrayLike, *names: AxisName) → NamedArray

Wraps a positional array as a NamedArray.

Parameters:

• 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.