def order_as(self, *axis_order: AxisName) -> NamedArray:
    """Ensures that the named axes are stored in this order, keeping them named.

    This function can be used if it is important for the axis names to appear in
    a consistent order, e.g. to ensure that two `NamedArray` instances have
    exactly the same PyTree structure.

    If you want a canonical ordering for a named array that doesn't involve
    knowing all the axis names in advance, you could do something like
    ``array.order_as(*sorted(array.named_shape.keys()))``.

    See also `order_like`.

    Args:
      *axis_order: Axis names in the order they should appear in the data array.
        Must be a permutation of all of the axis names in this array.

    Returns:
      Equivalent `NamedArray` whose data array contains the positional axes
      followed by the named axes in the given order.
    """
    # Create temporary "names" for the positional axes by creating new objects,
    # which are hashable but only compare equal by ID and are thus
    # guaranteed unique.
    tmp_names = [TmpPosAxisMarker() for _ in self.positional_shape]

    data_array = self.tag(*tmp_names).untag(*tmp_names, *axis_order).unwrap()
    return (
        NamedArray.wrap(data_array)
        .tag(*tmp_names, *axis_order)
        .untag(*tmp_names)
        .with_positional_prefix()
    )
Ensures that the named axes are stored in this order, keeping them named.

    This function can be used if it is important for the axis names to appear in
    a consistent order, e.g. to ensure that two `NamedArray` instances have
    exactly the same PyTree structure.

    If you want a canonical ordering for a named array that doesn't involve
    knowing all the axis names in advance, you could do something like
    ``array.order_as(*sorted(array.named_shape.keys()))``.

    See also `order_like`.

    Args:
      *axis_order: Axis names in the order they should appear in the data array.
        Must be a permutation of all of the axis names in this array.

    Returns:
      Equivalent `NamedArray` whose data array contains the positional axes
      followed by the named axes in the given order.
    
penzai/core/named_axes.py