scikit-hep / awkward

Manipulate JSON-like data with NumPy-like idioms.
https://awkward-array.org
BSD 3-Clause "New" or "Revised" License
848 stars 89 forks source link

ak.where unexpectedly fails on typetracers of option types #3321

Closed lgray closed 19 hours ago

lgray commented 20 hours ago

Version of Awkward Array

2.7.1

Description and code to reproduce

The following code:


import json
import awkward as ak

fromdict = {'class': 'RecordArray', 'fields': ['muon', 'jet'], 'contents': [{'class': 'ListOffsetArray', 'offsets': 'i64', 'content': {'class': 'RecordArray', 'fields': ['pt', 'eta', 'phi', 'crossref'], 'contents': [{'class': 'NumpyArray', 'pr\
imitive': 'int64', 'inner_shape': [], 'parameters': {}, 'form_key': "muon_pt!"}, {'class': 'NumpyArray', 'primitive': 'int64', 'inner_shape': [], 'parameters': {}, 'form_key': "muon_eta!"}, {'class': 'NumpyArray', 'primitive': 'int64', 'inner_\
shape': [], 'parameters': {}, 'form_key': "muon_phi!"}, {"class": "ListOffsetArray", "offsets": "i64", "content": {"class": "NumpyArray", "primitive": "int64", "inner_shape": [], "parameters": {}, "form_key": "muon_crossref_content!"}, "parame\
ters": {}, "form_key": "muon_crossref_index!"}], 'parameters': {}, 'form_key': "muon_record!"}, 'parameters': {}, 'form_key': "muon_list!"}, {'class': 'ListOffsetArray', 'offsets': 'i64', 'content': {'class': 'RecordArray', 'fields': ['pt', 'e\
ta', 'phi', "crossref", "thing1"], 'contents': [{'class': 'NumpyArray', 'primitive': 'int64', 'inner_shape': [], 'parameters': {}, 'form_key': "jet_pt!"}, {'class': 'NumpyArray', 'primitive': 'int64', 'inner_shape': [], 'parameters': {}, 'form\
_key': "jet_eta!"}, {'class': 'NumpyArray', 'primitive': 'int64', 'inner_shape': [], 'parameters': {}, 'form_key': "jet_phi!"}, {"class": "ListOffsetArray", "offsets": "i64", "content": {"class": "NumpyArray", "primitive": "int64", "inner_shap\
e": [], "parameters": {}, "form_key": "jet_crossref_content!"}, "parameters": {}, "form_key": "jet_crossref_index!"}, {'class': 'NumpyArray', 'primitive': 'int64', 'inner_shape': [], 'parameters': {}, 'form_key': "jet_thing1!"}], 'parameters':\
 {}, 'form_key': "jet_record!"}, 'parameters': {}, 'form_key': "jet_list!"}], 'parameters': {}, 'form_key': "outer!"}

form = ak.forms.from_dict(fromdict)

ttlayout, report = ak.typetracer.typetracer_with_report(form)

ttarray = ak.Array(ttlayout)

maybe = ak.firsts(ttarray)

varmaybe = ak.pad_none(ttarray, 3)

where = ak.where(abs(ttarray.jet.eta) < 1., 0.000511, ttarray.jet.thing1)

varmaybe_where = ak.where(abs(varmaybe.jet.eta) < 1., 0.000511, varmaybe.jet.thing1)

maybe_where = ak.where(abs(maybe.jet.eta) < 1., 0.000511, maybe.jet.thing1)

Fails with:

Traceback (most recent call last):
  File "/Users/lgray/coffea-dev/coffea/form_madness_again.py", line 18, in <module>
    varmaybe_where = ak.where(abs(varmaybe.jet.eta) < 1., 0.000511, varmaybe.jet.thing1)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_dispatch.py", line 64, in dispatch
    next(gen_or_result)
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/operations/ak_where.py", line 65, in where
    return _impl3(condition, x, y, mergebool, highlevel, behavior, attrs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/operations/ak_where.py", line 130, in _impl3
    out = ak._broadcasting.broadcast_and_apply(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1200, in broadcast_and_apply
    out = apply_step(
          ^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1178, in apply_step
    return continuation()
           ^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1147, in continuation
    return broadcast_any_list()
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_broadcasting.py", line 671, in broadcast_any_list
    outcontent = apply_step(
                 ^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1178, in apply_step
    return continuation()
           ^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1147, in continuation
    return broadcast_any_list()
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_broadcasting.py", line 747, in broadcast_any_list
    outcontent = apply_step(
                 ^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1178, in apply_step
    return continuation()
           ^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_broadcasting.py", line 1141, in continuation
    return broadcast_any_option_akwhere()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_broadcasting.py", line 863, in broadcast_any_option_akwhere
    NumpyArray(backend.nplike.zeros(len(inputs[2]), dtype=np.int8))
                                    ^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/contents/content.py", line 272, in __len__
    return int(self.length)
           ^^^^^^^^^^^^^^^^
  File "/Users/lgray/miniforge3/envs/coffea-dev/lib/python3.12/site-packages/awkward/_nplikes/shape.py", line 79, in __int__
    raise TypeError("cannot interpret unknown lengths as concrete values")
TypeError: cannot interpret unknown lengths as concrete values

This error occurred while calling

    ak.where(
        <Array-typetracer [...] type='## * var * ?bool'>
        0.000511
        <Array-typetracer [...] type='## * var * ?int64'>
    )
lgray commented 20 hours ago

Also: both jagged and flat code paths fail with the same message here. I have included both in the reproducer to make sure that all paths work after the fix.