V2: Overriding behavior in Numba

[Saransh-cpp]

Vector currently overrides the Numba behavior according to the documentation present here, but when switching to awkward._v2, the tests (or using the Awkward-Numba backend in any manner) gives the following error -

    def test():
        @numba.njit
        def extract(x):
            return x[2][0]

        array = vector.Array([[{"x": 1, "y": 2}], [], [{"x": 3, "y": 4}, {"x": 5, "y": 6}]])
>       out = extract(array)

tests/backends/test_awkward_numba.py:24:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/home/saransh/.local/lib/python3.9/site-packages/numba/core/dispatcher.py:468: in _compile_for_args
    error_rewrite(e, 'typing')
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

e = TypingError('Failed in nopython mode pipeline (step: nopython frontend)\nnon-precise type pyobject\nDuring: typing of ...lowing argument(s):\n- argument 0: Cannot determine Numba type of <class \'vector.backends.awkward.VectorArray2D\'>\n')
issue_type = 'typing'

    def error_rewrite(e, issue_type):
        """
        Rewrite and raise Exception `e` with help supplied based on the
        specified issue_type.
        """
        if config.SHOW_HELP:
            help_msg = errors.error_extras[issue_type]
            e.patch_message('\n'.join((str(e).rstrip(), help_msg)))
        if config.FULL_TRACEBACKS:
            raise e
        else:
>           raise e.with_traceback(None)
E           numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
E           non-precise type pyobject
E           During: typing of argument at /mnt/d/OpenSource/HEP/vector/tests/backends/test_awkward_numba.py (21)
E
E           File "tests/backends/test_awkward_numba.py", line 21:
E               def extract(x):
E                   return x[2][0]
E                   ^
E
E           This error may have been caused by the following argument(s):
E           - argument 0: Cannot determine Numba type of <class 'vector.backends.awkward.VectorArray2D'>

/home/saransh/.local/lib/python3.9/site-packages/numba/core/dispatcher.py:409: TypingError
-------------------------------------------------- Captured log call ---------------------------------------------------
DEBUG    numba.core.entrypoints:entrypoints.py:38 Loading extension: EntryPoint(name='init', value='awkward.numba:register', group='numba_extensions')
DEBUG    numba.core.entrypoints:entrypoints.py:46 Extension loading failed for: EntryPoint(name='init', value='awkward.numba:register', group='numba_extensions')
DEBUG    numba.core.entrypoints:entrypoints.py:38 Loading extension: EntryPoint(name='init', value='vector:register_numba', group='numba_extensions')
DEBUG    numba.core.byteflow:byteflow.py:74 bytecode dump:
>          0    NOP(arg=None, lineno=21)
           2    LOAD_FAST(arg=0, lineno=21)
           4    LOAD_CONST(arg=1, lineno=21)
           6    BINARY_SUBSCR(arg=None, lineno=21)
           8    LOAD_CONST(arg=2, lineno=21)
          10    BINARY_SUBSCR(arg=None, lineno=21)
          12    RETURN_VALUE(arg=None, lineno=21)
DEBUG    numba.core.byteflow:byteflow.py:100 pending: deque([State(pc_initial=0 nstack_initial=0)])
DEBUG    numba.core.byteflow:byteflow.py:103 stack: []
DEBUG    numba.core.byteflow:byteflow.py:278 dispatch pc=0, inst=NOP(arg=None, lineno=21)
DEBUG    numba.core.byteflow:byteflow.py:279 stack []
DEBUG    numba.core.byteflow:byteflow.py:278 dispatch pc=2, inst=LOAD_FAST(arg=0, lineno=21)
DEBUG    numba.core.byteflow:byteflow.py:279 stack []
DEBUG    numba.core.byteflow:byteflow.py:278 dispatch pc=4, inst=LOAD_CONST(arg=1, lineno=21)
DEBUG    numba.core.byteflow:byteflow.py:279 stack ['$x2.0']
DEBUG    numba.core.byteflow:byteflow.py:278 dispatch pc=6, inst=BINARY_SUBSCR(arg=None, lineno=21)
DEBUG    numba.core.byteflow:byteflow.py:279 stack ['$x2.0', '$const4.1']
DEBUG    numba.core.byteflow:byteflow.py:278 dispatch pc=8, inst=LOAD_CONST(arg=2, lineno=21)
DEBUG    numba.core.byteflow:byteflow.py:279 stack ['$6binary_subscr.2']
DEBUG    numba.core.byteflow:byteflow.py:278 dispatch pc=10, inst=BINARY_SUBSCR(arg=None, lineno=21)
DEBUG    numba.core.byteflow:byteflow.py:279 stack ['$6binary_subscr.2', '$const8.3']
DEBUG    numba.core.byteflow:byteflow.py:278 dispatch pc=12, inst=RETURN_VALUE(arg=None, lineno=21)
DEBUG    numba.core.byteflow:byteflow.py:279 stack ['$10binary_subscr.4']
DEBUG    numba.core.byteflow:byteflow.py:138 end state. edges=[]
DEBUG    numba.core.byteflow:byteflow.py:166 -------------------------Prune PHIs-------------------------
DEBUG    numba.core.byteflow:byteflow.py:235 Used_phis: defaultdict(<class 'set'>, {State(pc_initial=0 nstack_initial=0): set()})
DEBUG    numba.core.byteflow:byteflow.py:192 defmap: {}
DEBUG    numba.core.byteflow:byteflow.py:193 phismap: defaultdict(<class 'set'>, {})
DEBUG    numba.core.byteflow:byteflow.py:214 changing phismap: defaultdict(<class 'set'>, {})
DEBUG    numba.core.byteflow:byteflow.py:223 keep phismap: {}
DEBUG    numba.core.byteflow:byteflow.py:229 new_out: defaultdict(<class 'dict'>, {})
DEBUG    numba.core.byteflow:byteflow.py:239 ----------------------DONE Prune PHIs-----------------------
DEBUG    numba.core.byteflow:byteflow.py:150 block_infos State(pc_initial=0 nstack_initial=0):
AdaptBlockInfo(insts=((0, {}), (2, {'res': '$x2.0'}), (4, {'res': '$const4.1'}), (6, {'index': '$const4.1', 'target': '$x2.0', 'res': '$6binary_subscr.2'}), (8, {'res': '$const8.3'}), (10, {'index': '$const8.3', 'target': '$6binary_subscr.2', 'res': '$10binary_subscr.4'}), (12, {'retval': '$10binary_subscr.4', 'castval': '$12return_value.5'})), outgoing_phis={}, blockstack=(), active_try_block=None, outgoing_edgepushed={})
DEBUG    numba.core.interpreter:interpreter.py:1342 label 0:
    x = arg(0, name=x)                       ['x']
    $const4.1 = const(int, 2)                ['$const4.1']
    $6binary_subscr.2 = getitem(value=x, index=$const4.1, fn=<built-in function getitem>) ['$6binary_subscr.2', '$const4.1', 'x']
    $const8.3 = const(int, 0)                ['$const8.3']
    $10binary_subscr.4 = getitem(value=$6binary_subscr.2, index=$const8.3, fn=<built-in function getitem>) ['$10binary_subscr.4', '$6binary_subscr.2', '$const8.3']
    $12return_value.5 = cast(value=$10binary_subscr.4) ['$10binary_subscr.4', '$12return_value.5']
    return $12return_value.5                 ['$12return_value.5']

DEBUG    numba.core.ssa:ssa.py:162 ==== SSA block analysis pass on 0
DEBUG    numba.core.ssa:ssa.py:191 Running <numba.core.ssa._GatherDefsHandler object at 0x7f08be9159a0>
DEBUG    numba.core.ssa:ssa.py:193 on stmt: x = arg(0, name=x)
DEBUG    numba.core.ssa:ssa.py:193 on stmt: $const4.1 = const(int, 2)
DEBUG    numba.core.ssa:ssa.py:193 on stmt: $6binary_subscr.2 = static_getitem(value=x, index=2, index_var=$const4.1, fn=<built-in function getitem>)
DEBUG    numba.core.ssa:ssa.py:193 on stmt: $const8.3 = const(int, 0)
DEBUG    numba.core.ssa:ssa.py:193 on stmt: $10binary_subscr.4 = static_getitem(value=$6binary_subscr.2, index=0, index_var=$const8.3, fn=<built-in function getitem>)
DEBUG    numba.core.ssa:ssa.py:193 on stmt: $12return_value.5 = cast(value=$10binary_subscr.4)
DEBUG    numba.core.ssa:ssa.py:193 on stmt: return $12return_value.5
DEBUG    numba.core.ssa:ssa.py:154 defs defaultdict(<class 'list'>,
            {'$10binary_subscr.4': [<numba.core.ir.Assign object at 0x7f08be8b7fa0>],
             '$12return_value.5': [<numba.core.ir.Assign object at 0x7f08be8b7580>],
             '$6binary_subscr.2': [<numba.core.ir.Assign object at 0x7f08be8b7d90>],
             '$const4.1': [<numba.core.ir.Assign object at 0x7f08be915f40>],
             '$const8.3': [<numba.core.ir.Assign object at 0x7f08be8b7280>],
             'x': [<numba.core.ir.Assign object at 0x7f08be915ca0>]})
DEBUG    numba.core.ssa:ssa.py:156 SSA violators set()
DEBUG    numba.core.typeinfer:typeinfer.py:159 captured error
Traceback (most recent call last):
  File "/home/saransh/.local/lib/python3.9/site-packages/numba/core/typeinfer.py", line 155, in propagate
    constraint(typeinfer)
  File "/home/saransh/.local/lib/python3.9/site-packages/numba/core/typeinfer.py", line 229, in __call__
    raise TypingError('non-precise type {}'.format(ty))
numba.core.errors.TypingError: non-precise type pyobject
During: typing of argument at /mnt/d/OpenSource/HEP/vector/tests/backends/test_awkward_numba.py (21)
DEBUG    numba.core.typeinfer:typeinfer.py:159 captured error
Traceback (most recent call last):
  File "/home/saransh/.local/lib/python3.9/site-packages/numba/core/typeinfer.py", line 155, in propagate
    constraint(typeinfer)
  File "/home/saransh/.local/lib/python3.9/site-packages/numba/core/typeinfer.py", line 229, in __call__
    raise TypingError('non-precise type {}'.format(ty))
numba.core.errors.TypingError: non-precise type pyobject
During: typing of argument at /mnt/d/OpenSource/HEP/vector/tests/backends/test_awkward_numba.py (21)

I am not sure how this Numba behavior is supposed to be overridden in awkward._v2. Vector currently does this to override the behavior.

Additionally, I am not sure if this is a bug or a change in awkward's API, as I couldn't find any documentation on this. This brings me to my last question - Is there any deployed documentation for awkward._v2?

Thank you!

[agoose77]

Whilst we're in the v1→v2 transition, we haven't registered the NumPy or Numba dispatch mechanisms, so calling np.func(array) or jitting with Numba doesn't work out of the box. You can register the Numba machinery using the same code from v1 in ak._connect._numba.__init__:

@numba.extending.typeof_impl.register(ak.highlevel.Array)
def typeof_Array(obj, c):
    return obj.numba_type

@numba.extending.typeof_impl.register(ak.highlevel.Record)
def typeof_Record(obj, c):
    return obj.numba_type

@numba.extending.typeof_impl.register(ak.highlevel.ArrayBuilder)
def typeof_ArrayBuilder(obj, c):
    return obj.numba_type

[Saransh-cpp]

Thank you for the answer, @agoose77! This makes sense!