def into_pytree(self) -> jax.ShapeDtypeStruct | named_axes.NamedArray:
    """Converts an ``ArraySpec`` into a (possibly wrapped) PyTree leaf.

    By default, an ``ArraySpec`` has no PyTree children. This method can be
    used to convert it into a PyTree subtree that contains one leaf, a
    `jax.ShapeDtypeStruct`. This can then be used to e.g. restore parameters
    into
    a structure.

    If this structure has a named shape, will return a NamedArray wrapping a
    ``ShapeDtypeStruct``, with the parameter order inferred from the order of
    names in this ``ArraySpec``. Otherwise, will return an ordinary
    ``ShapeDtypeStruct``.

    Returns:
      A PyTree whose structure matches this structure.
    """
    if any(isinstance(dim, DimVar | MultiDimVar) for dim in self.shape) or any(
        isinstance(v, DimVar | RemainingAxisPlaceholder)
        or isinstance(k, MultiDimVar)
        for k, v in self.named_shape.items()
    ):
      raise ValueError(
          "Cannot use `into_pytree` on an ArraySpec with variable dimensions."
      )
    if self.named_shape:
      return named_axes.NamedArray(
          named_axes=collections.OrderedDict(self.named_shape.items()),
          data_array=jax.ShapeDtypeStruct(
              shape=self.shape + tuple(self.named_shape.values()),
              dtype=self.dtype,
          ),
      )
    else:
      return jax.ShapeDtypeStruct(self.shape, self.dtype)
Converts an ``ArraySpec`` into a (possibly wrapped) PyTree leaf.

    By default, an ``ArraySpec`` has no PyTree children. This method can be
    used to convert it into a PyTree subtree that contains one leaf, a
    `jax.ShapeDtypeStruct`. This can then be used to e.g. restore parameters
    into
    a structure.

    If this structure has a named shape, will return a NamedArray wrapping a
    ``ShapeDtypeStruct``, with the parameter order inferred from the order of
    names in this ``ArraySpec``. Otherwise, will return an ordinary
    ``ShapeDtypeStruct``.

    Returns:
      A PyTree whose structure matches this structure.
    
penzai/core/shapecheck.py