We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
unflatten raises counts must be an integer or a one-dimensional array of integers on a 1D awkward-array (worked in 1.x)
### Version of Awkward Array

2.0.5

### Description and code to reproduce

I'm trying to transition from awkward 1.x -> 2.x, and ran into an issue with `ak.unflatten` where the call in 1.x works, but fails in 2.x. In my code, I call:

```python
    num_constituents = ak.num(jets.constituents, axis=2)
    # Convert to offsets
    offsets_constituents = np.cumsum(np.asarray(ak.flatten(num_constituents, axis=1)))
    offsets_constituents = np.insert(offsets_constituents, 0, 0)
    # Then into starts and stops
    starts_constituents = ak.unflatten(offsets_constituents[:-1], ak.num(num_constituents, axis=1))
    stops_constituents = ak.unflatten(offsets_constituents[1:], ak.num(num_constituents, axis=1))
```

For a bit of context, I'm trying to convert some info out from c++ back to python, and have found it convenient to return flattened np arrays and then unflatten on return. In any case, `ak.num(num_constituents, axis=1)` is 1D, as far as I can tell / recall, but in awkward 2.x, I get the following ValueError:

```python
E               ValueError: while calling
E
E                   ak.unflatten(
E                       array = numpy.ndarray([  0   6  37 ... 397 422 442])
E                       counts = <Array [1, 1, 1, 1, 1, ..., 1, 1, 1, 2, 1] type='17 * int64'>
E                       axis = 0
E                       highlevel = True
E                       behavior = None
E                   )
E
E               Error details: counts must be an integer or a one-dimensional array of integers
```

With details on the variables here:

```
array      = array([  0,   6,  37,  50,  81, 118, 146, 173, 193, 208, 242, 275, 318,
       336, 364, 397, 422, 442])
axis       = 0
backend    = <awkward._backends.NumpyBackend object at 0x1109822f0>
behavior   = None
counts     = <IndexedArray len='17'>
    <index><Index dtype='int64' len='17'>
        [ 0  1  3  4  5  6  7  8  9 10 11 12 13 15 16 18 19]
    </Index></index>
    <content><NumpyArray dtype='int64' len='20'>
        [1 1 0 1 1 1 1 1 1 1 1 1 1 1 0 1 1 0 2 1]
    </NumpyArray></content>
</IndexedArray>
highlevel  = True
layout     = <NumpyArray dtype='int64' len='18'>
    [  0   6  37  50  81 118 146 173 193 208 242 275 318 336 364 397 422 442]
</NumpyArray>
```

If I wrap the `ak.num(num_constituents, axis=1)` call with `np.asarray`, my integration tests pass, which at least suggests that it works (although the tests definitely aren't comprehensive), so it may be somewhat trivial. If it just needs to be cast, I think it would be a nicer user experience if `unflatten` did this on it's own. Thanks!


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/awkward/operations/ak_unflatten.py`
Content:
```
1 # BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
2 
3 import numbers
4 
5 import awkward as ak
6 
7 np = ak._nplikes.NumpyMetadata.instance()
8 
9 
10 def unflatten(array, counts, axis=0, *, highlevel=True, behavior=None):
11     """
12     Args:
13         array: Array-like data (anything #ak.to_layout recognizes).
14         counts (int or array): Number of elements the new level should have.
15             If an integer, the new level will be regularly sized; otherwise,
16             it will consist of variable-length lists with the given lengths.
17         axis (int): The dimension at which this operation is applied. The
18             outermost dimension is `0`, followed by `1`, etc., and negative
19             values count backward from the innermost: `-1` is the innermost
20             dimension, `-2` is the next level up, etc.
21         highlevel (bool): If True, return an #ak.Array; otherwise, return
22             a low-level #ak.contents.Content subclass.
23         behavior (None or dict): Custom #ak.behavior for the output array, if
24             high-level.
25 
26     Returns an array with an additional level of nesting. This is roughly the
27     inverse of #ak.flatten, where `counts` were obtained by #ak.num (both with
28     `axis=1`).
29 
30     For example,
31 
32         >>> original = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]])
33         >>> counts = ak.num(original)
34         >>> array = ak.flatten(original)
35         >>> counts
36         <Array [3, 0, 2, 1, 4] type='5 * int64'>
37         >>> array
38         <Array [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] type='10 * int64'>
39         >>> ak.unflatten(array, counts)
40         <Array [[0, 1, 2], [], [3, ...], [5], [6, 7, 8, 9]] type='5 * var * int64'>
41 
42     An inner dimension can be unflattened by setting the `axis` parameter, but
43     operations like this constrain the `counts` more tightly.
44 
45     For example, we can subdivide an already divided list:
46 
47         >>> original = ak.Array([[1, 2, 3, 4], [], [5, 6, 7], [8, 9]])
48         >>> ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).show()
49         [[[1, 2], [3, 4]],
50          [],
51          [[5], [6, 7]],
52          [[8], [9]]]
53 
54     But the counts have to add up to the lengths of those lists. We can't mix
55     values from the first `[1, 2, 3, 4]` with values from the next `[5, 6, 7]`.
56 
57         >>> ak.unflatten(original, [2, 1, 2, 2, 1, 1], axis=1).show()
58         ValueError: while calling
59             ak.unflatten(
60                 array = <Array [[1, 2, 3, 4], [], ..., [8, 9]] type='4 * var * int64'>
61                 counts = [2, 1, 2, 2, 1, 1]
62                 axis = 1
63                 highlevel = True
64                 behavior = None
65             )
66         Error details: structure imposed by 'counts' does not fit in the array or partition at axis=1
67 
68     Also note that new lists created by this function cannot cross partitions
69     (which is only possible at `axis=0`, anyway).
70 
71     See also #ak.num and #ak.flatten.
72     """
73     with ak._errors.OperationErrorContext(
74         "ak.unflatten",
75         dict(
76             array=array,
77             counts=counts,
78             axis=axis,
79             highlevel=highlevel,
80             behavior=behavior,
81         ),
82     ):
83         return _impl(array, counts, axis, highlevel, behavior)
84 
85 
86 def _impl(array, counts, axis, highlevel, behavior):
87     layout = ak.operations.to_layout(
88         array, allow_record=False, allow_other=False
89     ).to_packed()
90     behavior = ak._util.behavior_of(array, behavior=behavior)
91     backend = layout.backend
92 
93     if isinstance(counts, (numbers.Integral, np.integer)):
94         current_offsets = None
95     else:
96         counts = ak.operations.to_layout(
97             counts, allow_record=False, allow_other=False
98         ).to_packed()
99         if counts.is_option and (counts.content.is_numpy or counts.content.is_unknown):
100             mask = counts.mask_as_bool(valid_when=False)
101             counts = backend.nplike.to_rectilinear(
102                 ak.operations.fill_none(counts, 0, axis=-1, highlevel=False)
103             )
104         elif counts.is_numpy or counts.is_unknown:
105             counts = backend.nplike.to_rectilinear(counts)
106             mask = False
107         else:
108             raise ak._errors.wrap_error(
109                 ValueError(
110                     "counts must be an integer or a one-dimensional array of integers"
111                 )
112             )
113 
114         if counts.ndim != 1:
115             raise ak._errors.wrap_error(ValueError("counts must be one-dimensional"))
116 
117         if not issubclass(counts.dtype.type, np.integer):
118             raise ak._errors.wrap_error(ValueError("counts must be integers"))
119 
120         current_offsets = backend.index_nplike.empty(len(counts) + 1, np.int64)
121         current_offsets[0] = 0
122         backend.index_nplike.cumsum(counts, out=current_offsets[1:])
123 
124     def unflatten_this_layout(layout):
125         nonlocal current_offsets
126 
127         if isinstance(counts, (numbers.Integral, np.integer)):
128             if counts < 0 or counts > len(layout):
129                 raise ak._errors.wrap_error(
130                     ValueError("too large counts for array or negative counts")
131                 )
132             out = ak.contents.RegularArray(layout, counts)
133 
134         else:
135             position = (
136                 backend.index_nplike.searchsorted(
137                     current_offsets,
138                     backend.index_nplike.array([len(layout)]),
139                     side="right",
140                 )[0]
141                 - 1
142             )
143             if position >= len(current_offsets) or current_offsets[position] != len(
144                 layout
145             ):
146                 raise ak._errors.wrap_error(
147                     ValueError(
148                         "structure imposed by 'counts' does not fit in the array or partition "
149                         "at axis={}".format(axis)
150                     )
151                 )
152 
153             offsets = current_offsets[: position + 1]
154             current_offsets = current_offsets[position:] - len(layout)
155 
156             out = ak.contents.ListOffsetArray(ak.index.Index64(offsets), layout)
157             if not isinstance(mask, (bool, np.bool_)):
158                 index = ak.index.Index8(
159                     backend.nplike.asarray(mask).astype(np.int8), nplike=backend.nplike
160                 )
161                 out = ak.contents.ByteMaskedArray(index, out, valid_when=False)
162 
163         return out
164 
165     if axis == 0 or ak._util.maybe_posaxis(layout, axis, 1) == 0:
166         out = unflatten_this_layout(layout)
167 
168     else:
169 
170         def apply(layout, depth, **kwargs):
171             posaxis = ak._util.maybe_posaxis(layout, axis, depth)
172             if posaxis == depth and layout.is_list:
173                 # We are one *above* the level where we want to apply this.
174                 listoffsetarray = layout.to_ListOffsetArray64(True)
175                 outeroffsets = backend.index_nplike.asarray(listoffsetarray.offsets)
176 
177                 content = unflatten_this_layout(
178                     listoffsetarray.content[: outeroffsets[-1]]
179                 )
180                 if isinstance(content, ak.contents.ByteMaskedArray):
181                     inneroffsets = backend.index_nplike.asarray(content.content.offsets)
182                 elif isinstance(content, ak.contents.RegularArray):
183                     inneroffsets = backend.index_nplike.asarray(
184                         content.to_ListOffsetArray64(True).offsets
185                     )
186                 else:
187                     inneroffsets = backend.index_nplike.asarray(content.offsets)
188 
189                 positions = (
190                     backend.index_nplike.searchsorted(
191                         inneroffsets, outeroffsets, side="right"
192                     )
193                     - 1
194                 )
195                 if not backend.index_nplike.array_equal(
196                     inneroffsets[positions], outeroffsets
197                 ):
198                     raise ak._errors.wrap_error(
199                         ValueError(
200                             "structure imposed by 'counts' does not fit in the array or partition "
201                             "at axis={}".format(axis)
202                         )
203                     )
204                 positions[0] = 0
205 
206                 return ak.contents.ListOffsetArray(ak.index.Index64(positions), content)
207 
208         out = ak._do.recursively_apply(layout, apply)
209 
210     if current_offsets is not None and not (
211         len(current_offsets) == 1 and current_offsets[0] == 0
212     ):
213         raise ak._errors.wrap_error(
214             ValueError(
215                 "structure imposed by 'counts' does not fit in the array or partition "
216                 "at axis={}".format(axis)
217             )
218         )
219 
220     return ak._util.wrap(out, behavior, highlevel)
221
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/awkward/operations/ak_unflatten.py b/src/awkward/operations/ak_unflatten.py
--- a/src/awkward/operations/ak_unflatten.py
+++ b/src/awkward/operations/ak_unflatten.py
@@ -93,9 +93,10 @@
     if isinstance(counts, (numbers.Integral, np.integer)):
         current_offsets = None
     else:
-        counts = ak.operations.to_layout(
-            counts, allow_record=False, allow_other=False
-        ).to_packed()
+        counts = ak.operations.to_layout(counts, allow_record=False, allow_other=False)
+        if counts.is_indexed and not counts.is_option:
+            counts = counts.project()
+
         if counts.is_option and (counts.content.is_numpy or counts.content.is_unknown):
             mask = counts.mask_as_bool(valid_when=False)
             counts = backend.nplike.to_rectilinear(

{"golden_diff": "diff --git a/src/awkward/operations/ak_unflatten.py b/src/awkward/operations/ak_unflatten.py\n--- a/src/awkward/operations/ak_unflatten.py\n+++ b/src/awkward/operations/ak_unflatten.py\n@@ -93,9 +93,10 @@\n     if isinstance(counts, (numbers.Integral, np.integer)):\n         current_offsets = None\n     else:\n-        counts = ak.operations.to_layout(\n-            counts, allow_record=False, allow_other=False\n-        ).to_packed()\n+        counts = ak.operations.to_layout(counts, allow_record=False, allow_other=False)\n+        if counts.is_indexed and not counts.is_option:\n+            counts = counts.project()\n+\n         if counts.is_option and (counts.content.is_numpy or counts.content.is_unknown):\n             mask = counts.mask_as_bool(valid_when=False)\n             counts = backend.nplike.to_rectilinear(\n", "issue": "unflatten raises counts must be an integer or a one-dimensional array of integers on a 1D awkward-array (worked in 1.x)\n### Version of Awkward Array\n\n2.0.5\n\n### Description and code to reproduce\n\nI'm trying to transition from awkward 1.x -> 2.x, and ran into an issue with `ak.unflatten` where the call in 1.x works, but fails in 2.x. In my code, I call:\r\n\r\n```python\r\n    num_constituents = ak.num(jets.constituents, axis=2)\r\n    # Convert to offsets\r\n    offsets_constituents = np.cumsum(np.asarray(ak.flatten(num_constituents, axis=1)))\r\n    offsets_constituents = np.insert(offsets_constituents, 0, 0)\r\n    # Then into starts and stops\r\n    starts_constituents = ak.unflatten(offsets_constituents[:-1], ak.num(num_constituents, axis=1))\r\n    stops_constituents = ak.unflatten(offsets_constituents[1:], ak.num(num_constituents, axis=1))\r\n```\r\n\r\nFor a bit of context, I'm trying to convert some info out from c++ back to python, and have found it convenient to return flattened np arrays and then unflatten on return. In any case, `ak.num(num_constituents, axis=1)` is 1D, as far as I can tell / recall, but in awkward 2.x, I get the following ValueError:\r\n\r\n```python\r\nE               ValueError: while calling\r\nE\r\nE                   ak.unflatten(\r\nE                       array = numpy.ndarray([  0   6  37 ... 397 422 442])\r\nE                       counts = <Array [1, 1, 1, 1, 1, ..., 1, 1, 1, 2, 1] type='17 * int64'>\r\nE                       axis = 0\r\nE                       highlevel = True\r\nE                       behavior = None\r\nE                   )\r\nE\r\nE               Error details: counts must be an integer or a one-dimensional array of integers\r\n```\r\n\r\nWith details on the variables here:\r\n\r\n```\r\narray      = array([  0,   6,  37,  50,  81, 118, 146, 173, 193, 208, 242, 275, 318,\r\n       336, 364, 397, 422, 442])\r\naxis       = 0\r\nbackend    = <awkward._backends.NumpyBackend object at 0x1109822f0>\r\nbehavior   = None\r\ncounts     = <IndexedArray len='17'>\r\n    <index><Index dtype='int64' len='17'>\r\n        [ 0  1  3  4  5  6  7  8  9 10 11 12 13 15 16 18 19]\r\n    </Index></index>\r\n    <content><NumpyArray dtype='int64' len='20'>\r\n        [1 1 0 1 1 1 1 1 1 1 1 1 1 1 0 1 1 0 2 1]\r\n    </NumpyArray></content>\r\n</IndexedArray>\r\nhighlevel  = True\r\nlayout     = <NumpyArray dtype='int64' len='18'>\r\n    [  0   6  37  50  81 118 146 173 193 208 242 275 318 336 364 397 422 442]\r\n</NumpyArray>\r\n```\r\n\r\nIf I wrap the `ak.num(num_constituents, axis=1)` call with `np.asarray`, my integration tests pass, which at least suggests that it works (although the tests definitely aren't comprehensive), so it may be somewhat trivial. If it just needs to be cast, I think it would be a nicer user experience if `unflatten` did this on it's own. Thanks!\r\n\n", "before_files": [{"content": "# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE\n\nimport numbers\n\nimport awkward as ak\n\nnp = ak._nplikes.NumpyMetadata.instance()\n\n\ndef unflatten(array, counts, axis=0, *, highlevel=True, behavior=None):\n    \"\"\"\n    Args:\n        array: Array-like data (anything #ak.to_layout recognizes).\n        counts (int or array): Number of elements the new level should have.\n            If an integer, the new level will be regularly sized; otherwise,\n            it will consist of variable-length lists with the given lengths.\n        axis (int): The dimension at which this operation is applied. The\n            outermost dimension is `0`, followed by `1`, etc., and negative\n            values count backward from the innermost: `-1` is the innermost\n            dimension, `-2` is the next level up, etc.\n        highlevel (bool): If True, return an #ak.Array; otherwise, return\n            a low-level #ak.contents.Content subclass.\n        behavior (None or dict): Custom #ak.behavior for the output array, if\n            high-level.\n\n    Returns an array with an additional level of nesting. This is roughly the\n    inverse of #ak.flatten, where `counts` were obtained by #ak.num (both with\n    `axis=1`).\n\n    For example,\n\n        >>> original = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]])\n        >>> counts = ak.num(original)\n        >>> array = ak.flatten(original)\n        >>> counts\n        <Array [3, 0, 2, 1, 4] type='5 * int64'>\n        >>> array\n        <Array [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] type='10 * int64'>\n        >>> ak.unflatten(array, counts)\n        <Array [[0, 1, 2], [], [3, ...], [5], [6, 7, 8, 9]] type='5 * var * int64'>\n\n    An inner dimension can be unflattened by setting the `axis` parameter, but\n    operations like this constrain the `counts` more tightly.\n\n    For example, we can subdivide an already divided list:\n\n        >>> original = ak.Array([[1, 2, 3, 4], [], [5, 6, 7], [8, 9]])\n        >>> ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).show()\n        [[[1, 2], [3, 4]],\n         [],\n         [[5], [6, 7]],\n         [[8], [9]]]\n\n    But the counts have to add up to the lengths of those lists. We can't mix\n    values from the first `[1, 2, 3, 4]` with values from the next `[5, 6, 7]`.\n\n        >>> ak.unflatten(original, [2, 1, 2, 2, 1, 1], axis=1).show()\n        ValueError: while calling\n            ak.unflatten(\n                array = <Array [[1, 2, 3, 4], [], ..., [8, 9]] type='4 * var * int64'>\n                counts = [2, 1, 2, 2, 1, 1]\n                axis = 1\n                highlevel = True\n                behavior = None\n            )\n        Error details: structure imposed by 'counts' does not fit in the array or partition at axis=1\n\n    Also note that new lists created by this function cannot cross partitions\n    (which is only possible at `axis=0`, anyway).\n\n    See also #ak.num and #ak.flatten.\n    \"\"\"\n    with ak._errors.OperationErrorContext(\n        \"ak.unflatten\",\n        dict(\n            array=array,\n            counts=counts,\n            axis=axis,\n            highlevel=highlevel,\n            behavior=behavior,\n        ),\n    ):\n        return _impl(array, counts, axis, highlevel, behavior)\n\n\ndef _impl(array, counts, axis, highlevel, behavior):\n    layout = ak.operations.to_layout(\n        array, allow_record=False, allow_other=False\n    ).to_packed()\n    behavior = ak._util.behavior_of(array, behavior=behavior)\n    backend = layout.backend\n\n    if isinstance(counts, (numbers.Integral, np.integer)):\n        current_offsets = None\n    else:\n        counts = ak.operations.to_layout(\n            counts, allow_record=False, allow_other=False\n        ).to_packed()\n        if counts.is_option and (counts.content.is_numpy or counts.content.is_unknown):\n            mask = counts.mask_as_bool(valid_when=False)\n            counts = backend.nplike.to_rectilinear(\n                ak.operations.fill_none(counts, 0, axis=-1, highlevel=False)\n            )\n        elif counts.is_numpy or counts.is_unknown:\n            counts = backend.nplike.to_rectilinear(counts)\n            mask = False\n        else:\n            raise ak._errors.wrap_error(\n                ValueError(\n                    \"counts must be an integer or a one-dimensional array of integers\"\n                )\n            )\n\n        if counts.ndim != 1:\n            raise ak._errors.wrap_error(ValueError(\"counts must be one-dimensional\"))\n\n        if not issubclass(counts.dtype.type, np.integer):\n            raise ak._errors.wrap_error(ValueError(\"counts must be integers\"))\n\n        current_offsets = backend.index_nplike.empty(len(counts) + 1, np.int64)\n        current_offsets[0] = 0\n        backend.index_nplike.cumsum(counts, out=current_offsets[1:])\n\n    def unflatten_this_layout(layout):\n        nonlocal current_offsets\n\n        if isinstance(counts, (numbers.Integral, np.integer)):\n            if counts < 0 or counts > len(layout):\n                raise ak._errors.wrap_error(\n                    ValueError(\"too large counts for array or negative counts\")\n                )\n            out = ak.contents.RegularArray(layout, counts)\n\n        else:\n            position = (\n                backend.index_nplike.searchsorted(\n                    current_offsets,\n                    backend.index_nplike.array([len(layout)]),\n                    side=\"right\",\n                )[0]\n                - 1\n            )\n            if position >= len(current_offsets) or current_offsets[position] != len(\n                layout\n            ):\n                raise ak._errors.wrap_error(\n                    ValueError(\n                        \"structure imposed by 'counts' does not fit in the array or partition \"\n                        \"at axis={}\".format(axis)\n                    )\n                )\n\n            offsets = current_offsets[: position + 1]\n            current_offsets = current_offsets[position:] - len(layout)\n\n            out = ak.contents.ListOffsetArray(ak.index.Index64(offsets), layout)\n            if not isinstance(mask, (bool, np.bool_)):\n                index = ak.index.Index8(\n                    backend.nplike.asarray(mask).astype(np.int8), nplike=backend.nplike\n                )\n                out = ak.contents.ByteMaskedArray(index, out, valid_when=False)\n\n        return out\n\n    if axis == 0 or ak._util.maybe_posaxis(layout, axis, 1) == 0:\n        out = unflatten_this_layout(layout)\n\n    else:\n\n        def apply(layout, depth, **kwargs):\n            posaxis = ak._util.maybe_posaxis(layout, axis, depth)\n            if posaxis == depth and layout.is_list:\n                # We are one *above* the level where we want to apply this.\n                listoffsetarray = layout.to_ListOffsetArray64(True)\n                outeroffsets = backend.index_nplike.asarray(listoffsetarray.offsets)\n\n                content = unflatten_this_layout(\n                    listoffsetarray.content[: outeroffsets[-1]]\n                )\n                if isinstance(content, ak.contents.ByteMaskedArray):\n                    inneroffsets = backend.index_nplike.asarray(content.content.offsets)\n                elif isinstance(content, ak.contents.RegularArray):\n                    inneroffsets = backend.index_nplike.asarray(\n                        content.to_ListOffsetArray64(True).offsets\n                    )\n                else:\n                    inneroffsets = backend.index_nplike.asarray(content.offsets)\n\n                positions = (\n                    backend.index_nplike.searchsorted(\n                        inneroffsets, outeroffsets, side=\"right\"\n                    )\n                    - 1\n                )\n                if not backend.index_nplike.array_equal(\n                    inneroffsets[positions], outeroffsets\n                ):\n                    raise ak._errors.wrap_error(\n                        ValueError(\n                            \"structure imposed by 'counts' does not fit in the array or partition \"\n                            \"at axis={}\".format(axis)\n                        )\n                    )\n                positions[0] = 0\n\n                return ak.contents.ListOffsetArray(ak.index.Index64(positions), content)\n\n        out = ak._do.recursively_apply(layout, apply)\n\n    if current_offsets is not None and not (\n        len(current_offsets) == 1 and current_offsets[0] == 0\n    ):\n        raise ak._errors.wrap_error(\n            ValueError(\n                \"structure imposed by 'counts' does not fit in the array or partition \"\n                \"at axis={}\".format(axis)\n            )\n        )\n\n    return ak._util.wrap(out, behavior, highlevel)\n", "path": "src/awkward/operations/ak_unflatten.py"}], "after_files": [{"content": "# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE\n\nimport numbers\n\nimport awkward as ak\n\nnp = ak._nplikes.NumpyMetadata.instance()\n\n\ndef unflatten(array, counts, axis=0, *, highlevel=True, behavior=None):\n    \"\"\"\n    Args:\n        array: Array-like data (anything #ak.to_layout recognizes).\n        counts (int or array): Number of elements the new level should have.\n            If an integer, the new level will be regularly sized; otherwise,\n            it will consist of variable-length lists with the given lengths.\n        axis (int): The dimension at which this operation is applied. The\n            outermost dimension is `0`, followed by `1`, etc., and negative\n            values count backward from the innermost: `-1` is the innermost\n            dimension, `-2` is the next level up, etc.\n        highlevel (bool): If True, return an #ak.Array; otherwise, return\n            a low-level #ak.contents.Content subclass.\n        behavior (None or dict): Custom #ak.behavior for the output array, if\n            high-level.\n\n    Returns an array with an additional level of nesting. This is roughly the\n    inverse of #ak.flatten, where `counts` were obtained by #ak.num (both with\n    `axis=1`).\n\n    For example,\n\n        >>> original = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]])\n        >>> counts = ak.num(original)\n        >>> array = ak.flatten(original)\n        >>> counts\n        <Array [3, 0, 2, 1, 4] type='5 * int64'>\n        >>> array\n        <Array [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] type='10 * int64'>\n        >>> ak.unflatten(array, counts)\n        <Array [[0, 1, 2], [], [3, ...], [5], [6, 7, 8, 9]] type='5 * var * int64'>\n\n    An inner dimension can be unflattened by setting the `axis` parameter, but\n    operations like this constrain the `counts` more tightly.\n\n    For example, we can subdivide an already divided list:\n\n        >>> original = ak.Array([[1, 2, 3, 4], [], [5, 6, 7], [8, 9]])\n        >>> ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).show()\n        [[[1, 2], [3, 4]],\n         [],\n         [[5], [6, 7]],\n         [[8], [9]]]\n\n    But the counts have to add up to the lengths of those lists. We can't mix\n    values from the first `[1, 2, 3, 4]` with values from the next `[5, 6, 7]`.\n\n        >>> ak.unflatten(original, [2, 1, 2, 2, 1, 1], axis=1).show()\n        ValueError: while calling\n            ak.unflatten(\n                array = <Array [[1, 2, 3, 4], [], ..., [8, 9]] type='4 * var * int64'>\n                counts = [2, 1, 2, 2, 1, 1]\n                axis = 1\n                highlevel = True\n                behavior = None\n            )\n        Error details: structure imposed by 'counts' does not fit in the array or partition at axis=1\n\n    Also note that new lists created by this function cannot cross partitions\n    (which is only possible at `axis=0`, anyway).\n\n    See also #ak.num and #ak.flatten.\n    \"\"\"\n    with ak._errors.OperationErrorContext(\n        \"ak.unflatten\",\n        dict(\n            array=array,\n            counts=counts,\n            axis=axis,\n            highlevel=highlevel,\n            behavior=behavior,\n        ),\n    ):\n        return _impl(array, counts, axis, highlevel, behavior)\n\n\ndef _impl(array, counts, axis, highlevel, behavior):\n    layout = ak.operations.to_layout(\n        array, allow_record=False, allow_other=False\n    ).to_packed()\n    behavior = ak._util.behavior_of(array, behavior=behavior)\n    backend = layout.backend\n\n    if isinstance(counts, (numbers.Integral, np.integer)):\n        current_offsets = None\n    else:\n        counts = ak.operations.to_layout(counts, allow_record=False, allow_other=False)\n        if counts.is_indexed and not counts.is_option:\n            counts = counts.project()\n\n        if counts.is_option and (counts.content.is_numpy or counts.content.is_unknown):\n            mask = counts.mask_as_bool(valid_when=False)\n            counts = backend.nplike.to_rectilinear(\n                ak.operations.fill_none(counts, 0, axis=-1, highlevel=False)\n            )\n        elif counts.is_numpy or counts.is_unknown:\n            counts = backend.nplike.to_rectilinear(counts)\n            mask = False\n        else:\n            raise ak._errors.wrap_error(\n                ValueError(\n                    \"counts must be an integer or a one-dimensional array of integers\"\n                )\n            )\n\n        if counts.ndim != 1:\n            raise ak._errors.wrap_error(ValueError(\"counts must be one-dimensional\"))\n\n        if not issubclass(counts.dtype.type, np.integer):\n            raise ak._errors.wrap_error(ValueError(\"counts must be integers\"))\n\n        current_offsets = backend.index_nplike.empty(len(counts) + 1, np.int64)\n        current_offsets[0] = 0\n        backend.index_nplike.cumsum(counts, out=current_offsets[1:])\n\n    def unflatten_this_layout(layout):\n        nonlocal current_offsets\n\n        if isinstance(counts, (numbers.Integral, np.integer)):\n            if counts < 0 or counts > len(layout):\n                raise ak._errors.wrap_error(\n                    ValueError(\"too large counts for array or negative counts\")\n                )\n            out = ak.contents.RegularArray(layout, counts)\n\n        else:\n            position = (\n                backend.index_nplike.searchsorted(\n                    current_offsets,\n                    backend.index_nplike.array([len(layout)]),\n                    side=\"right\",\n                )[0]\n                - 1\n            )\n            if position >= len(current_offsets) or current_offsets[position] != len(\n                layout\n            ):\n                raise ak._errors.wrap_error(\n                    ValueError(\n                        \"structure imposed by 'counts' does not fit in the array or partition \"\n                        \"at axis={}\".format(axis)\n                    )\n                )\n\n            offsets = current_offsets[: position + 1]\n            current_offsets = current_offsets[position:] - len(layout)\n\n            out = ak.contents.ListOffsetArray(ak.index.Index64(offsets), layout)\n            if not isinstance(mask, (bool, np.bool_)):\n                index = ak.index.Index8(\n                    backend.nplike.asarray(mask).astype(np.int8), nplike=backend.nplike\n                )\n                out = ak.contents.ByteMaskedArray(index, out, valid_when=False)\n\n        return out\n\n    if axis == 0 or ak._util.maybe_posaxis(layout, axis, 1) == 0:\n        out = unflatten_this_layout(layout)\n\n    else:\n\n        def apply(layout, depth, **kwargs):\n            posaxis = ak._util.maybe_posaxis(layout, axis, depth)\n            if posaxis == depth and layout.is_list:\n                # We are one *above* the level where we want to apply this.\n                listoffsetarray = layout.to_ListOffsetArray64(True)\n                outeroffsets = backend.index_nplike.asarray(listoffsetarray.offsets)\n\n                content = unflatten_this_layout(\n                    listoffsetarray.content[: outeroffsets[-1]]\n                )\n                if isinstance(content, ak.contents.ByteMaskedArray):\n                    inneroffsets = backend.index_nplike.asarray(content.content.offsets)\n                elif isinstance(content, ak.contents.RegularArray):\n                    inneroffsets = backend.index_nplike.asarray(\n                        content.to_ListOffsetArray64(True).offsets\n                    )\n                else:\n                    inneroffsets = backend.index_nplike.asarray(content.offsets)\n\n                positions = (\n                    backend.index_nplike.searchsorted(\n                        inneroffsets, outeroffsets, side=\"right\"\n                    )\n                    - 1\n                )\n                if not backend.index_nplike.array_equal(\n                    inneroffsets[positions], outeroffsets\n                ):\n                    raise ak._errors.wrap_error(\n                        ValueError(\n                            \"structure imposed by 'counts' does not fit in the array or partition \"\n                            \"at axis={}\".format(axis)\n                        )\n                    )\n                positions[0] = 0\n\n                return ak.contents.ListOffsetArray(ak.index.Index64(positions), content)\n\n        out = ak._do.recursively_apply(layout, apply)\n\n    if current_offsets is not None and not (\n        len(current_offsets) == 1 and current_offsets[0] == 0\n    ):\n        raise ak._errors.wrap_error(\n            ValueError(\n                \"structure imposed by 'counts' does not fit in the array or partition \"\n                \"at axis={}\".format(axis)\n            )\n        )\n\n    return ak._util.wrap(out, behavior, highlevel)\n", "path": "src/awkward/operations/ak_unflatten.py"}]}