fix: support typetracer in ak.unflatten

[agoose77]

This PR fixes assumptions of length and known values that dask-awkward's test-suite fails.

[codecov]

Codecov Report

Merging #2293 (18548cd) into main (7fc01d8) will increase coverage by 0.08%.
The diff coverage is 100.00%.

Additional details and impacted files

Impacted Files	Coverage Δ
src/awkward/typing.py	88.88% <ø> (+5.55%)	⬆️
src/awkward/_util.py	75.60% <100.00%> (+0.59%)	⬆️
src/awkward/contents/listoffsetarray.py	80.65% <100.00%> (-0.06%)	⬇️
src/awkward/operations/ak_unflatten.py	95.83% <100.00%> (+0.31%)	⬆️
src/awkward/operations/ak_pad_none.py	100.00% <0.00%> (ø)
src/awkward/operations/ak_drop_none.py	100.00% <0.00%> (ø)
src/awkward/operations/ak_local_index.py	100.00% <0.00%> (ø)
...rc/awkward/operations/ak_merge_union_of_records.py	98.05% <0.00%> (+0.01%)	⬆️
src/awkward/operations/ak_concatenate.py	96.26% <0.00%> (+0.03%)	⬆️
src/awkward/operations/ak_cartesian.py	90.69% <0.00%> (+0.07%)	⬆️
... and 36 more

[jpivarski]

We talked about this today, and I thought that if threading a typetracer through the whole function is difficult, it could be a candidate for a specialized path. But it looks like you did it!

I don't see any tests, though. (Tests for the specific case of ak.unflatten with typetracers.) Do they come up in dask-awkward?

[jpivarski]

So... counts was originally obtained via to_layout (layouts have .length, not .size), but then it was replaced by to_backend_array. I assume that means it could be NumPy, CuPy, a JAX DeviceArray, or a TypeTracerArray. All of these have a .size?

For a NumPy array, the .size is the product of shape, not its first element. Is there any way that the NumpyArray might be multidimensional? I suppose not, if it's a valid counts.

[agoose77]

Yes, a previous PR explicitly ensures that counts is a backend array, all of which have .size. We have an explicit check on counts.ndim to ensure that it's 1D!

[agoose77]

I don't see any tests, though. (Tests for the specific case of ak.unflatten with typetracers.) Do they come up in dask-awkward?

Dask-awkward does test this, because it brought this to my attention from their test suite. However, we should really test it for ourselves (I'm being lazy 🤕). I'll add some tests before we merge.

[jpivarski]

✔️ on tests/test_2293_unflatten_typetracer.py

[agoose77]

Ah, I noticed a final corner-case: counts can be known for typetracer backends, and I'm trying to preserve that kind of information where possible. I modified the PR to normalise counts if its value is not known, rather than testing if the backend has known_data.