Repository: scikit-hep/awkward

Issue:
`_packed` negative axes aren't handled properly for unions
When `_packed` is given a negative axis for a layout that doesn't return a positive value in `wrap_axis_if_negative`, it actually bails early:
https://github.com/scikit-hep/awkward-1.0/blob/bec1a0d240f0e590a835a0ef70e3687624c96be5/src/awkward/operations/structure.py#L2122-L2123

What we are trying to catch here is whether the current depth is past the terminal depth, i.e. if we have gone past the deepest level we wish to operate on. This doesn't mean that deeper layouts aren't modified; the parent layout can simplify the child, but the child itself will not be simplified beyond this. 

Because this routine bails whenever `posaxis` is negative, it is "broken" for unions.

This bug is only triggered by the public API if we try to `unflatten` an array containing unions using a negative axis.


Discussion:
@jpivarski I think the fix for this should look like

```python3
    def is_below_axis(layout, depth, axis):
        if axis >= 0:
            # Depth starts at 1
            return depth-1 > axis
        else:
            # We only consider ourselves to have passed the union depth 
            # if we cannot process any union members, i.e.
            # if `abs(axis)` is farther than the farthest leaf
            min_depth, max_depth = layout.minmax_depth
            return max_depth + axis < 0
        
    def apply(layout, depth, posaxis):
        if posaxis is not None:
            # If a particular axis was given, do not proceed past that axis
            if is_below_axis(layout, depth, posaxis):
                return layout
 ```

This is fiddly stuff, does it look right to you? The only problem here is that we don't handle the case where the initial layout is not deep enough; `is_below_axis` (by design) returns `True` (and doesn't raise an `Exception`) for any axis ~ ∞. I could add a check outside of the recursion that axis does not exceed the `max_depth` of the outer layout.

Alternatively, we drop the `axis` parameter. I don't know whether this is a good idea - the fix doesn't seem complicated, and I think we might benefit from not over simplifying. But, perhaps simpler is better?
To explain the above modified code snippet, consider this example:
```python3
layout = ak.layout.ListOffsetArray64(
    ak.layout.Index64(np.r_[1, 3]),
    ak.layout.UnionArray8_64(
        ak.layout.Index8(np.r_[0, 1, 0, 1]),
        ak.layout.Index64(np.r_[0, 0, 1, 1]),
        [
            ak.layout.ListOffsetArray64(
                ak.layout.Index64(np.r_[1, 4, 8]), ak.layout.NumpyArray(np.arange(128))
            ),
            ak.layout.ListOffsetArray64(
                ak.layout.Index64(np.r_[1, 3, 5]),
                ak.layout.ListOffsetArray64(
                    ak.layout.Index64(np.r_[1, 3, 6, 9, 12, 16, 18, 25]),
                    ak.layout.NumpyArray(np.arange(64) - 32),
                ),
            ),
        ],
    ),
)
```
For `axis=2`, the effective depth is actually `depth=3`. When walking the tree, we wait until we hit that depth

* `ListOffsetArray64(depth=1)`
   * `UnionArray8_64(depth=2)` 
   * `ListOffsetArray64(depth=2)`
      * `NumpyArray(depth=3)`
   * `ListOffsetArray64(depth=2)`
      * `ListOffsetArray64(depth=3)`
         * ~~`NumpyArray(depth=4)`~~
      
and then stop processing. This means that the final `NumpyArray(depth=4)` is not processed.
For negative axes, the behaviour is different - we work backwards from the leaves. This means that for `axis=-2` we see:
* `ListOffsetArray64(depth=1, neg_depth=?)`
   * `UnionArray8_64(depth=2, neg_depth=?)` 
   * `ListOffsetArray64(depth=2, neg_depth=-2)`
      * ~~`NumpyArray(depth=3, neg_depth=-1)`~~
   * `ListOffsetArray64(depth=2, neg_depth=-3)`
      * `ListOffsetArray64(depth=3, neg_depth=-2)`
         * ~~`NumpyArray(depth=4, neg_depth=-1)`~~

i.e. we don't process the leaves at `neg_depth=-1`.
(I'm thinking this through; chances are I'm saying stuff you already know or have already figured out...)

`ak.packed` doesn't have an `axis` for simplicity—because of what the function does, I don't think anyone would ever want to pack only one dimension. So this is only about the internal `_packed`, which is used by `ak.unflatten`.

An `axis` shouldn't be considered valid on a UnionArray unless it can be used in some way by every `content` in that union. Your examples have a list-type in the first `content` and a list-list-type in the second `content`, so the second should be able to take one more possible negative `axis` than the first. That is, if the `axis` of the function refers to a numeric level, the first could take `axis=-1` and `axis=-2` and the second could take `axis=-1`, `axis=-2`, and `axis=-3`. If the function's `axis` refers to a point between dimensions, like where to put the new counts in `ak.unflatten`, then it would be just `axis=-1` for the first and `axis=-1` or `axis=-2` for the second.

The `depth` counts up from the bottom, so `depth` is either equal to the negation of a negative `axis` or off by one (consistently). Since `minmax_depth` gives you both the minimum and maximum and the UnionArray has to work for _all_ `content` branches, I think it's the _minimum_ that you want to compare against a negated `axis`, not the maximum, as in your code example. When you ask for `minmax_depth` of a UnionArray, the minimum is the one that's more restrictive in what negative axes are allowed.

I'll try drawing one of these diagrams, too:

```
                               min max depth   allowed axis?
------------------------------------------------------------
list-type                      3   4           0 (only)
│
└── union-type                 2   3           none!
    │
    ├── list-type              2   2           1, -2
    │   └── numeric-type       1   1           2, -1
    │
    └── list-type              3   3           1, -3
        └── list-type          2   2           2, -2
            └── numeric-type   1   1           3, -1
```

I'm confident about the `minmax_depth` values because I computed them in a terminal. That's just what you get when you call the function. I'm uncertain about the "allowed axis?" column—what I've written is my intuition about what should happen, but it could be different for different kinds of functions (some apply _at_ a dimension and others apply _above_ a dimension) and I couldn't find any examples of functions that work! UnionArrays break a lot of functionality. If there is some function that successfully passes an `axis` through a UnionArray, I don't know offhand which it is. RecordArrays and negative `axis` generally work.

But anyway, the "allowed axis?" column represents the theory: negative `axis` should count up from the bottom and as such, it can be different in different branches of a UnionArray or RecordArray (the only two layout types that can have more than one `content`). If it's ever ambiguous, it's not allowed, so a negative axis can't be applied above a split-point with different depths.

I'm starting to remember—it might have been an across-the-board decision to make negative `axis` never work through UnionArrays. Unions are different from records because in a RecordArray, every element of the array has all branches, but in a UnionArray, each array element has one of the possible branches. I think that the decision was that interpreting `axis` differently on an array element-by-array element basis was too subtle/not understandable enough to be a good policy. Maybe you shouldn't work so hard to get this one to accept negative `axis`.
Thanks for the detailed reply, the diagram is much clearer than mine. You're correct that this *does not* matter for `ak.packed`, only `ak.unflatten` (because we choose to optimise the case).

This entire issue is really about the fact that neither positive nor negative axes are well defined for layouts with unions containing differing numbers of list type layouts. I made a decision in the original implementation & this issue. In my mind, there are really two separate problems:

## What negative axes "mean" (in general)
I wasn't aware that we deliberately disallow ambiguous allowed axes. This makes sense, and changes the interpretation of `minmax`. I was using `max_depth` because we are using negative depths whilst traversing the tree from the root. In this way, I interpret the meaning of negative axes for `_packed` to be 

> pack each layout until we reach the given axis, and then stop. 

With this interpretation, we expect any layout which contains some `content` that has sufficient depth (`max_depth`) to be visited, which is why I use `max_depth` in the condition.

The question is whether this interpretation is valid. Effectively, unions with (`min != max`) contents require us to choose whether out-of-bounds axes (e.g. an `axis` which is too large for one union content, but is valid for another, or an `axis` which is too negative for one content, but not the other) are legal or not. In the case of `_packed`, I think this is an oversight of mine - I arbitrarily chose negative depths to be interpreted in terms of the deepest branch, but one could equally take the minimum. 
Perhaps it is better to consider only the minimum, in the same way that one treats the positive case (as you suggest).

## How `_packed` should be fixed
`_packed` takes an axis so that it can prepare an array for `unflatten`. This function is different to many of the other Awkward functions that take axes, because it operates *above* the given axis, which makes the above ambiguity a problem. 

What we *want* to happen is that each layout is packed as it is visited on our way to unflattening at a given depth. Then we defer the meaning of `axis` to that which is defined for `unflatten`. I think this is the right solution. Either we move the recursion from `_packed` into `packed`, and have `_packed` perform a single packing operation, or we replace `axis` with `depth`, which defines the recursion depth. I think the latter is actually going to be simpler than having to implement the recursion in `packed`.
What you propose sounds like the right solution, to the extent that you need to make `ak.unflatten` work not just for data with unions, but with different list-depth unions. ("Sounds like" to the degree that I've understood it.)

If you don't have an explicit need for this case and if the solution would be hard to understand even when it works, then it may not be worthwhile. If you do decide that it's not worthwhile, at least it should raise a good error message.
I think we're on the same page. Having taken a step back, unpacking above a depth really only makes sense for root-defined axes (aka depths) rather than allowing negative axes, so removing that feature isn't really a problem from a doing-it-the-right-way perspective.

Resolving patch:
diff --git a/src/awkward/_util.py b/src/awkward/_util.py
--- a/src/awkward/_util.py
+++ b/src/awkward/_util.py
@@ -1177,12 +1177,11 @@ def recursively_apply(
     getfunction,
     pass_depth=True,
     pass_user=False,
-    pass_apply=False,
     user=None,
     keep_parameters=True,
     numpy_to_regular=False,
 ):
-    def apply(layout, depth, user):
+    def transform(layout, depth, user):
         if numpy_to_regular and isinstance(layout, ak.layout.NumpyArray):
             layout = layout.toRegularArray()
 
@@ -1191,214 +1190,17 @@ def apply(layout, depth, user):
             args = args + (depth,)
         if pass_user:
             args = args + (user,)
-        if pass_apply:
-            args = args + (apply,)
 
         custom = getfunction(layout, *args)
         if callable(custom):
             return custom()
-        else:
-            user = custom
-
-        # the rest of this is one switch statement
-        if isinstance(layout, ak.partition.PartitionedArray):
-            return ak.partition.IrregularlyPartitionedArray(
-                [apply(x, depth, user) for x in layout.partitions]
-            )
-
-        elif isinstance(layout, ak.layout.NumpyArray):
-            if keep_parameters:
-                return layout
-            else:
-                return ak.layout.NumpyArray(
-                    ak.nplike.of(layout).asarray(layout), layout.identities, None
-                )
-
-        elif isinstance(layout, ak.layout.EmptyArray):
-            if keep_parameters:
-                return layout
-            else:
-                return ak.layout.EmptyArray(layout.identities, None)
-
-        elif isinstance(layout, ak.layout.RegularArray):
-            return ak.layout.RegularArray(
-                apply(layout.content, depth + 1, user),
-                layout.size,
-                len(layout),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.ListArray32):
-            return ak.layout.ListArray32(
-                layout.starts,
-                layout.stops,
-                apply(layout.content, depth + 1, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.ListArrayU32):
-            return ak.layout.ListArrayU32(
-                layout.starts,
-                layout.stops,
-                apply(layout.content, depth + 1, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.ListArray64):
-            return ak.layout.ListArray64(
-                layout.starts,
-                layout.stops,
-                apply(layout.content, depth + 1, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.ListOffsetArray32):
-            return ak.layout.ListOffsetArray32(
-                layout.offsets,
-                apply(layout.content, depth + 1, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.ListOffsetArrayU32):
-            return ak.layout.ListOffsetArrayU32(
-                layout.offsets,
-                apply(layout.content, depth + 1, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.ListOffsetArray64):
-            return ak.layout.ListOffsetArray64(
-                layout.offsets,
-                apply(layout.content, depth + 1, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.IndexedArray32):
-            return ak.layout.IndexedArray32(
-                layout.index,
-                apply(layout.content, depth, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.IndexedArrayU32):
-            return ak.layout.IndexedArrayU32(
-                layout.index,
-                apply(layout.content, depth, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.IndexedArray64):
-            return ak.layout.IndexedArray64(
-                layout.index,
-                apply(layout.content, depth, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.IndexedOptionArray32):
-            return ak.layout.IndexedOptionArray32(
-                layout.index,
-                apply(layout.content, depth, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.IndexedOptionArray64):
-            return ak.layout.IndexedOptionArray64(
-                layout.index,
-                apply(layout.content, depth, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.ByteMaskedArray):
-            return ak.layout.ByteMaskedArray(
-                layout.mask,
-                apply(layout.content, depth, user),
-                layout.valid_when,
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.BitMaskedArray):
-            return ak.layout.BitMaskedArray(
-                layout.mask,
-                apply(layout.content, depth, user),
-                layout.valid_when,
-                len(layout),
-                layout.lsb_order,
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.UnmaskedArray):
-            return ak.layout.UnmaskedArray(
-                apply(layout.content, depth, user),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.RecordArray):
-            return ak.layout.RecordArray(
-                [apply(x, depth, user) for x in layout.contents],
-                layout.recordlookup,
-                len(layout),
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.Record):
-            return ak.layout.Record(
-                apply(layout.array, depth, user),
-                layout.at,
-            )
-
-        elif isinstance(layout, ak.layout.UnionArray8_32):
-            return ak.layout.UnionArray8_32(
-                layout.tags,
-                layout.index,
-                [apply(x, depth, user) for x in layout.contents],
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.UnionArray8_U32):
-            return ak.layout.UnionArray8_U32(
-                layout.tags,
-                layout.index,
-                [apply(x, depth, user) for x in layout.contents],
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.UnionArray8_64):
-            return ak.layout.UnionArray8_64(
-                layout.tags,
-                layout.index,
-                [apply(x, depth, user) for x in layout.contents],
-                layout.identities,
-                layout.parameters if keep_parameters else None,
-            )
-
-        elif isinstance(layout, ak.layout.VirtualArray):
-            return apply(layout.array, depth, user)
 
         else:
-            raise AssertionError(
-                "unrecognized Content type: {0}".format(type(layout))
-                + exception_suffix(__file__)
+            return transform_child_layouts(
+                transform, layout, depth, user=custom, keep_parameters=keep_parameters
             )
 
-    return apply(layout, 1, user)
+    return transform(layout, 1, user)
 
 
 def recursive_walk(layout, apply, args=(), depth=1, materialize=False):
@@ -1469,6 +1271,206 @@ def recursive_walk(layout, apply, args=(), depth=1, materialize=False):
         )
 
 
+def transform_child_layouts(transform, layout, depth, user=None, keep_parameters=True):
+    # the rest of this is one switch statement
+    if isinstance(layout, ak.partition.PartitionedArray):
+        return ak.partition.IrregularlyPartitionedArray(
+            [transform(x, depth, user) for x in layout.partitions]
+        )
+
+    elif isinstance(layout, ak.layout.NumpyArray):
+        if keep_parameters:
+            return layout
+        else:
+            return ak.layout.NumpyArray(
+                ak.nplike.of(layout).asarray(layout), layout.identities, None
+            )
+
+    elif isinstance(layout, ak.layout.EmptyArray):
+        if keep_parameters:
+            return layout
+        else:
+            return ak.layout.EmptyArray(layout.identities, None)
+
+    elif isinstance(layout, ak.layout.RegularArray):
+        return ak.layout.RegularArray(
+            transform(layout.content, depth + 1, user),
+            layout.size,
+            len(layout),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.ListArray32):
+        return ak.layout.ListArray32(
+            layout.starts,
+            layout.stops,
+            transform(layout.content, depth + 1, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.ListArrayU32):
+        return ak.layout.ListArrayU32(
+            layout.starts,
+            layout.stops,
+            transform(layout.content, depth + 1, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.ListArray64):
+        return ak.layout.ListArray64(
+            layout.starts,
+            layout.stops,
+            transform(layout.content, depth + 1, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.ListOffsetArray32):
+        return ak.layout.ListOffsetArray32(
+            layout.offsets,
+            transform(layout.content, depth + 1, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.ListOffsetArrayU32):
+        return ak.layout.ListOffsetArrayU32(
+            layout.offsets,
+            transform(layout.content, depth + 1, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.ListOffsetArray64):
+        return ak.layout.ListOffsetArray64(
+            layout.offsets,
+            transform(layout.content, depth + 1, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.IndexedArray32):
+        return ak.layout.IndexedArray32(
+            layout.index,
+            transform(layout.content, depth, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.IndexedArrayU32):
+        return ak.layout.IndexedArrayU32(
+            layout.index,
+            transform(layout.content, depth, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.IndexedArray64):
+        return ak.layout.IndexedArray64(
+            layout.index,
+            transform(layout.content, depth, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.IndexedOptionArray32):
+        return ak.layout.IndexedOptionArray32(
+            layout.index,
+            transform(layout.content, depth, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.IndexedOptionArray64):
+        return ak.layout.IndexedOptionArray64(
+            layout.index,
+            transform(layout.content, depth, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.ByteMaskedArray):
+        return ak.layout.ByteMaskedArray(
+            layout.mask,
+            transform(layout.content, depth, user),
+            layout.valid_when,
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.BitMaskedArray):
+        return ak.layout.BitMaskedArray(
+            layout.mask,
+            transform(layout.content, depth, user),
+            layout.valid_when,
+            len(layout),
+            layout.lsb_order,
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.UnmaskedArray):
+        return ak.layout.UnmaskedArray(
+            transform(layout.content, depth, user),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.RecordArray):
+        return ak.layout.RecordArray(
+            [transform(x, depth, user) for x in layout.contents],
+            layout.recordlookup,
+            len(layout),
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.Record):
+        return ak.layout.Record(
+            transform(layout.array, depth, user),
+            layout.at,
+        )
+
+    elif isinstance(layout, ak.layout.UnionArray8_32):
+        return ak.layout.UnionArray8_32(
+            layout.tags,
+            layout.index,
+            [transform(x, depth, user) for x in layout.contents],
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.UnionArray8_U32):
+        return ak.layout.UnionArray8_U32(
+            layout.tags,
+            layout.index,
+            [transform(x, depth, user) for x in layout.contents],
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.UnionArray8_64):
+        return ak.layout.UnionArray8_64(
+            layout.tags,
+            layout.index,
+            [transform(x, depth, user) for x in layout.contents],
+            layout.identities,
+            layout.parameters if keep_parameters else None,
+        )
+
+    elif isinstance(layout, ak.layout.VirtualArray):
+        return transform(layout.array, depth, user)
+
+    else:
+        raise AssertionError(
+            "unrecognized Content type: {0}".format(type(layout))
+            + exception_suffix(__file__)
+        )
+
+
 def find_caches(layout):
     # Both of the implementations below find referentially unique mutablemappings,
     # but the PartitionedArray case is optimized for many unique values (with a set)
diff --git a/src/awkward/operations/structure.py b/src/awkward/operations/structure.py
--- a/src/awkward/operations/structure.py
+++ b/src/awkward/operations/structure.py
@@ -1958,10 +1958,6 @@ def unflatten(array, counts, axis=0, highlevel=True, behavior=None):
     layout = ak.operations.convert.to_layout(
         array, allow_record=False, allow_other=False
     )
-    # Pack the layout above the axis that will be unflattened. This ensures that
-    # the `counts` array, which is computed through these layouts, aligns with
-    # the layout to be unflattened (#910)
-    layout = _packed(layout, axis=axis - 1, highlevel=False)
 
     if isinstance(counts, (numbers.Integral, np.integer)):
         current_offsets = None
@@ -2039,7 +2035,12 @@ def doit(layout):
 
     else:
 
-        def getfunction(layout, depth, posaxis):
+        def transform(layout, depth, posaxis):
+            # Pack the current layout. This ensures that the `counts` array,
+            # which is computed with these layouts applied, aligns with the
+            # internal layout to be unflattened (#910)
+            layout = _pack_layout(layout)
+
             posaxis = layout.axis_wrap_if_negative(posaxis)
             if posaxis == depth and isinstance(layout, ak._util.listtypes):
                 # We are one *above* the level where we want to apply this.
@@ -2065,34 +2066,22 @@ def getfunction(layout, depth, posaxis):
                         "at axis={0}".format(axis) + ak._util.exception_suffix(__file__)
                     )
 
-                return lambda: ak.layout.ListOffsetArray64(
+                return ak.layout.ListOffsetArray64(
                     ak.layout.Index64(positions), content
                 )
 
             else:
-                return posaxis
+                return ak._util.transform_child_layouts(
+                    transform, layout, depth, posaxis
+                )
 
         if isinstance(layout, ak.partition.PartitionedArray):
             outparts = []
             for part in layout.partitions:
-                outparts.append(
-                    ak._util.recursively_apply(
-                        part,
-                        getfunction,
-                        pass_depth=True,
-                        pass_user=True,
-                        user=axis,
-                    )
-                )
+                outparts.append(transform(part, depth=1, posaxis=axis))
             out = ak.partition.IrregularlyPartitionedArray(outparts)
         else:
-            out = ak._util.recursively_apply(
-                layout,
-                getfunction,
-                pass_depth=True,
-                pass_user=True,
-                user=axis,
-            )
+            out = transform(layout, depth=1, posaxis=axis)
 
     if current_offsets is not None and not (
         len(current_offsets[0]) == 1 and current_offsets[0][0] == 0
@@ -2105,201 +2094,171 @@ def getfunction(layout, depth, posaxis):
     return ak._util.maybe_wrap_like(out, array, behavior, highlevel)
 
 
-def _packed(array, axis=None, highlevel=True, behavior=None):
-    layout = ak.operations.convert.to_layout(
-        array, allow_record=True, allow_other=False
-    )
+def _pack_layout(layout):
     nplike = ak.nplike.of(layout)
 
-    def truncate(layout, n):
-        return layout[:n]
+    if isinstance(layout, ak.layout.NumpyArray):
+        return layout.contiguous()
 
-    def apply(layout, depth, posaxis):
-        if posaxis is not None:
-            # If a particular axis was given
-            posaxis = layout.axis_wrap_if_negative(posaxis)
-            # Do not proceed past that axis
-            if posaxis < depth - 1:
-                return layout
+    # EmptyArray is a no-op
+    elif isinstance(layout, ak.layout.EmptyArray):
+        return layout
 
-        if isinstance(layout, ak.layout.NumpyArray):
-            return layout.contiguous()
+    # Project indexed arrays
+    elif isinstance(layout, ak._util.indexedoptiontypes):
+        if isinstance(layout.content, ak._util.optiontypes):
+            return layout.simplify()
 
-        # EmptyArray is a no-op
-        elif isinstance(layout, ak.layout.EmptyArray):
-            return layout
+        index = nplike.asarray(layout.index)
+        new_index = nplike.zeros_like(index)
 
-        # Project indexed arrays
-        elif isinstance(layout, ak._util.indexedoptiontypes):
-            if isinstance(layout.content, ak._util.optiontypes):
-                return apply(layout.simplify(), depth, posaxis)
+        is_none = index < 0
+        new_index[is_none] = -1
+        new_index[~is_none] = nplike.arange(len(new_index) - nplike.sum(is_none))
 
-            index = nplike.asarray(layout.index)
-            new_index = nplike.zeros_like(index)
-
-            is_none = index < 0
-            new_index[is_none] = -1
-            new_index[~is_none] = nplike.arange(len(new_index) - nplike.sum(is_none))
-
-            return ak.layout.IndexedOptionArray64(
-                ak.layout.Index64(new_index), apply(layout.project(), depth, posaxis)
-            )
-
-        # Project indexed arrays
-        elif isinstance(layout, ak._util.indexedtypes):
-            return apply(layout.project(), depth, posaxis)
-
-        # ListArray performs both ordering and resizing
-        elif isinstance(
-            layout,
-            (
-                ak.layout.ListArray32,
-                ak.layout.ListArrayU32,
-                ak.layout.ListArray64,
-            ),
-        ):
-            return apply(layout.toListOffsetArray64(True), depth, posaxis)
+        return ak.layout.IndexedOptionArray64(
+            ak.layout.Index64(new_index), layout.project()
+        )
 
-        # ListOffsetArray performs resizing
-        elif isinstance(
-            layout,
-            (
-                ak.layout.ListOffsetArray32,
-                ak.layout.ListOffsetArray64,
-                ak.layout.ListOffsetArrayU32,
-            ),
-        ):
-            new_layout = layout.toListOffsetArray64(True)
-            new_length = new_layout.offsets[-1]
-            return ak.layout.ListOffsetArray64(
-                new_layout.offsets,
-                apply(truncate(new_layout.content, new_length), depth + 1, posaxis),
-                new_layout.identities,
-                new_layout.parameters,
-            )
+    # Project indexed arrays
+    elif isinstance(layout, ak._util.indexedtypes):
+        return layout.project()
 
-        # UnmaskedArray just wraps another array
-        elif isinstance(layout, ak.layout.UnmaskedArray):
-            return ak.layout.UnmaskedArray(apply(layout.content, depth, posaxis))
+    # ListArray performs both ordering and resizing
+    elif isinstance(
+        layout,
+        (
+            ak.layout.ListArray32,
+            ak.layout.ListArrayU32,
+            ak.layout.ListArray64,
+        ),
+    ):
+        return layout.toListOffsetArray64(True)
 
-        # UnionArrays can be simplified
-        # and their contents too
-        elif isinstance(layout, ak._util.uniontypes):
-            layout = layout.simplify()
+    # ListOffsetArray performs resizing
+    elif isinstance(
+        layout,
+        (
+            ak.layout.ListOffsetArray32,
+            ak.layout.ListOffsetArray64,
+            ak.layout.ListOffsetArrayU32,
+        ),
+    ):
+        new_layout = layout.toListOffsetArray64(True)
+        new_length = new_layout.offsets[-1]
+        return ak.layout.ListOffsetArray64(
+            new_layout.offsets,
+            new_layout.content[:new_length],
+            new_layout.identities,
+            new_layout.parameters,
+        )
 
-            # If we managed to lose the drop type entirely
-            if not isinstance(layout, ak._util.uniontypes):
-                return apply(layout, depth, posaxis)
+    # UnmaskedArray just wraps another array
+    elif isinstance(layout, ak.layout.UnmaskedArray):
+        return ak.layout.UnmaskedArray(layout.content)
 
-            # Pack simplified layout
-            tags = nplike.asarray(layout.tags)
-            index = nplike.asarray(layout.index)
+    # UnionArrays can be simplified
+    # and their contents too
+    elif isinstance(layout, ak._util.uniontypes):
+        layout = layout.simplify()
 
-            new_contents = [None] * len(layout.contents)
-            new_index = nplike.zeros_like(index)
+        # If we managed to lose the drop type entirely
+        if not isinstance(layout, ak._util.uniontypes):
+            return layout
 
-            # Compact indices
-            for i in range(len(layout.contents)):
-                is_i = tags == i
+        # Pack simplified layout
+        tags = nplike.asarray(layout.tags)
+        index = nplike.asarray(layout.index)
 
-                new_contents[i] = apply(layout.project(i), depth, posaxis)
-                new_index[is_i] = nplike.arange(nplike.sum(is_i))
+        new_contents = [None] * len(layout.contents)
+        new_index = nplike.zeros_like(index)
 
-            return ak.layout.UnionArray8_64(
-                ak.layout.Index8(tags),
-                ak.layout.Index64(new_index),
-                new_contents,
-                layout.identities,
-                layout.parameters,
-            )
+        # Compact indices
+        for i in range(len(layout.contents)):
+            is_i = tags == i
 
-        # RecordArray contents can be truncated
-        elif isinstance(layout, ak.layout.RecordArray):
-            return ak.layout.RecordArray(
-                [
-                    apply(truncate(c, len(layout)), depth, posaxis)
-                    for c in layout.contents
-                ],
-                layout.recordlookup,
-                len(layout),
-                layout.identities,
-                layout.parameters,
-            )
+            new_contents[i] = layout.project(i)
+            new_index[is_i] = nplike.arange(nplike.sum(is_i))
 
-        # RegularArrays can change length
-        elif isinstance(layout, ak.layout.RegularArray):
-            if not len(layout):
-                return layout
+        return ak.layout.UnionArray8_64(
+            ak.layout.Index8(tags),
+            ak.layout.Index64(new_index),
+            new_contents,
+            layout.identities,
+            layout.parameters,
+        )
 
-            content = layout.content
+    # RecordArray contents can be truncated
+    elif isinstance(layout, ak.layout.RecordArray):
+        return ak.layout.RecordArray(
+            [c[: len(layout)] for c in layout.contents],
+            layout.recordlookup,
+            len(layout),
+            layout.identities,
+            layout.parameters,
+        )
 
-            # Truncate content if it is larger than a perfect
-            # multiple of the RegularArray size
-            n, r = divmod(len(content), layout.size)
-            if r != 0:
-                content = truncate(content, n * layout.size)
+    # RegularArrays can change length
+    elif isinstance(layout, ak.layout.RegularArray):
+        if not len(layout):
+            return layout
 
-            return ak.layout.RegularArray(
-                apply(content, depth + 1, posaxis), layout.size
-            )
+        content = layout.content
 
-        # BitMaskedArrays can change length
-        elif isinstance(layout, ak.layout.BitMaskedArray):
-            layout = layout.simplify()
+        # Truncate content if it is larger than a perfect
+        # multiple of the RegularArray size
+        n, r = divmod(len(content), layout.size)
+        if r != 0:
+            content = content[: n * layout.size]
 
-            if not isinstance(ak.type(layout.content), ak.types.PrimitiveType):
-                return apply(layout.toIndexedOptionArray64(), depth, posaxis)
+        return ak.layout.RegularArray(content, layout.size)
 
-            return ak.layout.BitMaskedArray(
-                layout.mask,
-                apply(truncate(layout.content, len(layout)), depth, posaxis),
-                layout.valid_when,
-                len(layout),
-                layout.lsb_order,
-                layout.identities,
-                layout.parameters,
-            )
+    # BitMaskedArrays can change length
+    elif isinstance(layout, ak.layout.BitMaskedArray):
+        layout = layout.simplify()
 
-        # ByteMaskedArrays can change length
-        elif isinstance(layout, ak.layout.ByteMaskedArray):
-            layout = layout.simplify()
+        if not isinstance(ak.type(layout.content), ak.types.PrimitiveType):
+            return layout.toIndexedOptionArray64()
 
-            if not isinstance(ak.type(layout.content), ak.types.PrimitiveType):
-                return apply(layout.toIndexedOptionArray64(), depth, posaxis)
+        return ak.layout.BitMaskedArray(
+            layout.mask,
+            layout.content[: len(layout)],
+            layout.valid_when,
+            len(layout),
+            layout.lsb_order,
+            layout.identities,
+            layout.parameters,
+        )
 
-            return ak.layout.ByteMaskedArray(
-                layout.mask,
-                apply(truncate(layout.content, len(layout)), depth, posaxis),
-                layout.valid_when,
-                layout.identities,
-                layout.parameters,
-            )
+    # ByteMaskedArrays can change length
+    elif isinstance(layout, ak.layout.ByteMaskedArray):
+        layout = layout.simplify()
 
-        elif isinstance(layout, ak.layout.VirtualArray):
-            return apply(layout.array, depth, posaxis)
+        if not isinstance(ak.type(layout.content), ak.types.PrimitiveType):
+            return layout.toIndexedOptionArray64()
 
-        elif isinstance(layout, ak.partition.PartitionedArray):
-            return ak.partition.IrregularlyPartitionedArray(
-                [apply(x, depth, posaxis) for x in layout.partitions]
-            )
+        return ak.layout.ByteMaskedArray(
+            layout.mask,
+            layout.content[: len(layout)],
+            layout.valid_when,
+            layout.identities,
+            layout.parameters,
+        )
 
-        elif isinstance(layout, ak.layout.Record):
-            return ak.layout.Record(
-                apply(layout.array, depth, posaxis),
-                layout.at,
-            )
+    elif isinstance(layout, ak.layout.VirtualArray):
+        return layout.array
 
-        # Finally, fall through to failure
-        else:
-            raise AssertionError(
-                "unrecognized layout: "
-                + repr(layout)
-                + ak._util.exception_suffix(__file__)
-            )
+    elif isinstance(layout, ak.partition.PartitionedArray):
+        return layout
 
-    out = apply(layout, 1, axis)
+    elif isinstance(layout, ak.layout.Record):
+        return layout
 
-    return ak._util.maybe_wrap_like(out, array, behavior, highlevel)
+    # Finally, fall through to failure
+    else:
+        raise AssertionError(
+            "unrecognized layout: " + repr(layout) + ak._util.exception_suffix(__file__)
+        )
 
 
 def packed(array, highlevel=True, behavior=None):
@@ -2352,7 +2311,18 @@ def packed(array, highlevel=True, behavior=None):
 
     See also #ak.to_buffers.
     """
-    return _packed(array, highlevel=highlevel, behavior=behavior)
+    layout = ak.operations.convert.to_layout(
+        array, allow_record=True, allow_other=False
+    )
+
+    def transform(layout, depth=1, user=None):
+        return ak._util.transform_child_layouts(
+            transform, _pack_layout(layout), depth, user
+        )
+
+    out = transform(layout)
+
+    return ak._util.maybe_wrap_like(out, array, behavior, highlevel)
 
 
 def local_index(array, axis=-1, highlevel=True, behavior=None):


Test patch:
diff --git a/tests/test_0912-packed.py b/tests/test_0912-packed.py
--- a/tests/test_0912-packed.py
+++ b/tests/test_0912-packed.py
@@ -185,17 +185,3 @@ def test_record():
     packed = ak.packed(first, highlevel=False)
     assert ak.to_list(packed) == ak.to_list(first)
     assert len(packed.array) == len(layout)
-
-
-def test_axis():
-    content_0 = ak.layout.NumpyArray(np.arange(10))
-
-    offsets_1 = ak.layout.Index64(np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]))
-    content_1 = ak.layout.ListOffsetArray64(offsets_1, content_0)
-
-    layout = ak.layout.RegularArray(content_1, 3)
-
-    packed = ak.operations.structure._packed(layout, axis=0, highlevel=False)
-    assert ak.to_list(layout) == ak.to_list(packed)
-
-    assert np.asarray(packed.content.content).tolist() == np.asarray(content_0).tolist()