to_packed doesn't work on typetracer backed IndexOptionArray, BitMaskedArray, and UnionArray #3607
Description
Version of Awkward Array
main
Description and code to reproduce
To reproduce one example:
import awkward as ak
import numpy as np
index = np.array([0, 1, 2])
content = np.array([1, 2, 3])
layout = ak.contents.IndexedOptionArray(ak.index.Index(index), ak.contents.NumpyArray(content))
layout.to_typetracer().to_packed()Will error with
RuntimeError Traceback (most recent call last)
Cell In[4], line 1
----> 1 layout.to_typetracer().to_packed()
File ~/Dropbox/work/coffea_dev/awkward/src/awkward/contents/indexedoptionarray.py:1725, in IndexedOptionArray.to_packed(self, recursive)
1723 original_index = self._index.data
1724 is_none = original_index < 0
-> 1725 num_none = operator.index(nplike.count_nonzero(is_none))
1726 new_index = nplike.empty(self._index.length, dtype=self._index.dtype)
1727 if isinstance(nplike, Jax):
File ~/Dropbox/work/coffea_dev/awkward/src/awkward/_nplikes/typetracer.py:589, in TypeTracerArray.__index__(self)
588 def __index__(self) -> int:
--> 589 raise RuntimeError("cannot realise an unknown value")
RuntimeError: cannot realise an unknown value
The problem is that operator.index here:
Additionally there is another problem a few lines below here:
The problem there is in the typetracer case,
num_none is TypeTracerArray and new_index.size is either integer or unknown_length. new_index.size - num_none gives us a new TypeTracerArray which shape_item_as_index cannot handle and errors because TypeTracerArray is not a shape item basically.
Similarly, BitMaskedArray, and UnionArray don't properly handle typetracers in their to_packed implementations.