def order_like(
      self, other: NamedArray | NamedArrayView
  ) -> NamedArray | NamedArrayView:
    """Ensures that this array's PyTree structure matches another array's.

    This can be used to ensure that one named array has the same PyTree
    structure as another, so that the two can be jointly processed by
    non-namedarray-aware tree functions (e.g. `jax.tree_util` functions,
    `jax.lax.cond`, `jax.jvp`, etc).

    To ensure compatibility of entire PyTrees, you can use something like::

      jax.tree_util.tree_map(
          lambda a, b: a.order_like(b), tree1, tree2,
          is_leaf=pz.nx.is_namedarray,
      )

    Args:
      other: Another named array or named array view. Must have the same set of
        named axes as ``self``. If ``other`` is a `NamedArrayView`, ``other``
        must also have the same number of positional axes.

    Returns:
      A new `NamedArray` or `NamedArrayView` that has the content of ``self``
      but is possibly transposed to have the axes appear in the same order as
      ``other`` in the data array. If the arrays have the same named and
      positional shapes, the result will have the same PyTree structure as
      ``other``.
    """
    self.check_valid()
    other.check_valid()
    if isinstance(other, NamedArray):
      return self.order_as(*other.named_shape.keys())
    elif isinstance(other, NamedArrayView):
      if len(self.positional_shape) != len(other.positional_shape):
        raise ValueError(
            "Calling `order_like` with a NamedArrayView requires the two"
            " arrays to have the same number of positional axes, but got"
            f" positional shapes {self.positional_shape=},"
            f" {other.positional_shape=}"
        )
      if set(self.named_shape.keys()) != set(other.named_shape.keys()):
        raise ValueError(
            "Calling `order_like` with a NamedArrayView requires the two"
            " arrays to have the axis names, but got"
            f" named shapes {self.named_shape=}, {other.named_shape=}"
        )

      self_view = self.as_namedarrayview()
      new_to_old_data_axis = {}
      for old_data_axis, new_data_axis in zip(
          self_view.data_axis_for_logical_axis, other.data_axis_for_logical_axis
      ):
        new_to_old_data_axis[new_data_axis] = old_data_axis
      for name, new_data_axis in other.data_axis_for_name.items():
        new_to_old_data_axis[new_data_axis] = self_view.data_axis_for_name[name]
      new_data_array = jnp.transpose(
          self_view.data_array,
          [new_to_old_data_axis[i] for i in range(self_view.data_array.ndim)],
      )
      return NamedArrayView(
          data_shape=new_data_array.shape,
          data_axis_for_logical_axis=other.data_axis_for_logical_axis,
          data_axis_for_name=other.data_axis_for_name,
          data_array=new_data_array,
      )
    else:
      raise TypeError(
          "`order_like` requires a NamedArray or NamedArrayView, but got"
          f" {type(other).__name__}"
      )
Ensures that this array's PyTree structure matches another array's.

    This can be used to ensure that one named array has the same PyTree
    structure as another, so that the two can be jointly processed by
    non-namedarray-aware tree functions (e.g. `jax.tree_util` functions,
    `jax.lax.cond`, `jax.jvp`, etc).

    To ensure compatibility of entire PyTrees, you can use something like::

      jax.tree_util.tree_map(
          lambda a, b: a.order_like(b), tree1, tree2,
          is_leaf=pz.nx.is_namedarray,
      )

    Args:
      other: Another named array or named array view. Must have the same set of
        named axes as ``self``. If ``other`` is a `NamedArrayView`, ``other``
        must also have the same number of positional axes.

    Returns:
      A new `NamedArray` or `NamedArrayView` that has the content of ``self``
      but is possibly transposed to have the axes appear in the same order as
      ``other`` in the data array. If the arrays have the same named and
      positional shapes, the result will have the same PyTree structure as
      ``other``.
    
penzai/core/named_axes.py