Allow libraries that implement __array_ufunc to override DUFunc.c… #8995

jpivarski · 2023-05-31T22:30:51Z

This PR addresses a long-standing issue in which array-like objects that define an __array_ufunc__ method (NEP-18) can't be used with ufuncs created by np.vectorize unless a signature is provided. (I should have done this years ago!)

For instance, in Awkward Array,

>>> import awkward as ak
>>> import numba as nb

>>> nb.vectorize([nb.float64(nb.float64)])(lambda x: x + 1)(ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]]))
<Array [[2.1, 3.2, 4.3], [], [5.4, 6.5]] type='3 * var * float64'>

but

>>> nb.vectorize()(lambda x: x + 1)(ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]]))
<stdin>:1: NumbaWarning: 
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "<lambda>" failed type inference due to: non-precise type pyobject
During: typing of argument at <stdin> (1)

File "<stdin>", line 1:
<source missing, REPL/exec in use?>

/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/object_mode_passes.py:151: NumbaWarning: Function "<lambda>" was compiled in object mode without forceobj=True.

File "<stdin>", line 1:
<source missing, REPL/exec in use?>

  warnings.warn(errors.NumbaWarning(warn_msg,
/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/object_mode_passes.py:161: NumbaDeprecationWarning: 
Fall-back from the nopython compilation path to the object mode compilation path has been detected. This is deprecated behaviour that will be removed in Numba 0.59.0.

For more information visit https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit

File "<stdin>", line 1:
<source missing, REPL/exec in use?>

  warnings.warn(errors.NumbaDeprecationWarning(msg,
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/dufunc.py", line 206, in _compile_for_args
    return self._compile_for_argtys(tuple(argtys))
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/dufunc.py", line 225, in _compile_for_argtys
    actual_sig = ufuncbuilder._finalize_ufunc_signature(
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/ufuncbuilder.py", line 189, in _finalize_ufunc_signature
    raise TypeError("return type must be specified for object mode")
TypeError: return type must be specified for object mode

and even if we pass nopython=True,

>>> nb.vectorize(nopython=True)(lambda x: x + 1)(ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]]))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/dufunc.py", line 206, in _compile_for_args
    return self._compile_for_argtys(tuple(argtys))
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/dufunc.py", line 223, in _compile_for_argtys
    cres, argtys, return_type = ufuncbuilder._compile_element_wise_function(
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/ufuncbuilder.py", line 176, in _compile_element_wise_function
    cres = nb_func.compile(sig, **targetoptions)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/ufuncbuilder.py", line 124, in compile
    return self._compile_core(sig, flags, locals)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/ufuncbuilder.py", line 157, in _compile_core
    cres = compiler.compile_extra(typingctx, targetctx,
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler.py", line 742, in compile_extra
    return pipeline.compile_extra(func)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler.py", line 460, in compile_extra
    return self._compile_bytecode()
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler.py", line 528, in _compile_bytecode
    return self._compile_core()
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler.py", line 507, in _compile_core
    raise e
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler.py", line 494, in _compile_core
    pm.run(self.state)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 368, in run
    raise patched_exception
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 356, in run
    self._runPass(idx, pass_inst, state)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 311, in _runPass
    mutated |= check(pss.run_pass, internal_state)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 273, in check
    mangled = func(compiler_state)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/typed_passes.py", line 110, in run_pass
    typemap, return_type, calltypes, errs = type_inference_stage(
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/typed_passes.py", line 88, in type_inference_stage
    errs = infer.propagate(raise_errors=raise_errors)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/typeinfer.py", line 1086, in propagate
    raise errors[0]
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
non-precise type pyobject
During: typing of argument at <stdin> (1)

File "<stdin>", line 1:
<source missing, REPL/exec in use?>

Similarly, with the Sparse library,

>>> import numpy as np
>>> import sparse
>>> x = np.random.random((100, 100, 100))
>>> x[x < 0.9] = 0
>>> s = sparse.COO(x)

>>> nb.vectorize([nb.float64(nb.float64)], nopython=True)(lambda x: x + 100)(s)
<COO: shape=(100, 100, 100), dtype=float64, nnz=99450, fill_value=100.0>

>>> nb.vectorize([nb.float64(nb.float64)], nopython=True)(lambda x: x + 100)(s)[:2, :5, :8].todense()
array([[[100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.98842074],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.99039874, 100.        , 100.        ]],

       [[100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.99464796, 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.95169938, 100.        ]]])

but

>>> nb.vectorize()(lambda x: x + 100)(s)
<stdin>:1: NumbaWarning: 
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "<lambda>" failed type inference due to: non-precise type pyobject
During: typing of argument at <stdin> (1)

File "<stdin>", line 1:
<source missing, REPL/exec in use?>

/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/object_mode_passes.py:151: NumbaWarning: Function "<lambda>" was compiled in object mode without forceobj=True.

File "<stdin>", line 1:
<source missing, REPL/exec in use?>

  warnings.warn(errors.NumbaWarning(warn_msg,
/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/object_mode_passes.py:161: NumbaDeprecationWarning: 
Fall-back from the nopython compilation path to the object mode compilation path has been detected. This is deprecated behaviour that will be removed in Numba 0.59.0.

For more information visit https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit

File "<stdin>", line 1:
<source missing, REPL/exec in use?>

  warnings.warn(errors.NumbaDeprecationWarning(msg,
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/dufunc.py", line 206, in _compile_for_args
    return self._compile_for_argtys(tuple(argtys))
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/dufunc.py", line 225, in _compile_for_argtys
    actual_sig = ufuncbuilder._finalize_ufunc_signature(
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/ufuncbuilder.py", line 189, in _finalize_ufunc_signature
    raise TypeError("return type must be specified for object mode")
TypeError: return type must be specified for object mode

and

>>> nb.vectorize(nopython=True)(lambda x: x + 100)(s)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/dufunc.py", line 206, in _compile_for_args
    return self._compile_for_argtys(tuple(argtys))
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/dufunc.py", line 223, in _compile_for_argtys
    cres, argtys, return_type = ufuncbuilder._compile_element_wise_function(
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/ufuncbuilder.py", line 176, in _compile_element_wise_function
    cres = nb_func.compile(sig, **targetoptions)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/ufuncbuilder.py", line 124, in compile
    return self._compile_core(sig, flags, locals)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/np/ufunc/ufuncbuilder.py", line 157, in _compile_core
    cres = compiler.compile_extra(typingctx, targetctx,
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler.py", line 742, in compile_extra
    return pipeline.compile_extra(func)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler.py", line 460, in compile_extra
    return self._compile_bytecode()
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler.py", line 528, in _compile_bytecode
    return self._compile_core()
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler.py", line 507, in _compile_core
    raise e
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler.py", line 494, in _compile_core
    pm.run(self.state)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 368, in run
    raise patched_exception
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 356, in run
    self._runPass(idx, pass_inst, state)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 311, in _runPass
    mutated |= check(pss.run_pass, internal_state)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/compiler_machinery.py", line 273, in check
    mangled = func(compiler_state)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/typed_passes.py", line 110, in run_pass
    typemap, return_type, calltypes, errs = type_inference_stage(
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/typed_passes.py", line 88, in type_inference_stage
    errs = infer.propagate(raise_errors=raise_errors)
  File "/home/jpivarski/mambaforge/lib/python3.9/site-packages/numba/core/typeinfer.py", line 1086, in propagate
    raise errors[0]
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
non-precise type pyobject
During: typing of argument at <stdin> (1)

File "<stdin>", line 1:
<source missing, REPL/exec in use?>

However, with this PR, ufuncs created with @nb.vectorize and no type signature are recognized by Awkward Array:

>>> nb.vectorize()(lambda x: x + 1)(ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]]))
<Array [[2.1, 3.2, 4.3], [], [5.4, 6.5]] type='3 * var * float64'>

>>> nb.vectorize(nopython=True)(lambda x: x + 1)(ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]]))
<Array [[2.1, 3.2, 4.3], [], [5.4, 6.5]] type='3 * var * float64'>

and Sparse:

>>> nb.vectorize()(lambda x: x + 100)(s)
<COO: shape=(100, 100, 100), dtype=float64, nnz=100594, fill_value=100.0>

>>> nb.vectorize()(lambda x: x + 100)(s)[:2, :5, :8].todense()
array([[[100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.96245334],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ]],

       [[100.        , 100.        , 100.        , 100.        ,
         100.        , 100.99515858, 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.95034604, 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.96199577, 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ],
        [100.        , 100.        , 100.        , 100.        ,
         100.        , 100.        , 100.        , 100.        ]]])

It's because __array_ufunc__ methods are usually implemented by pre- and post-processing the result of applying the ufunc to some underlying NumPy array. If an __array_ufunc__ method is implemented some other way, without Numba-compatible objects at any level, then @nb.vectorize-decorated functions wouldn't work with it, anyway.

This is a draft PR because:

Please ensure that you have written units tests for the changes made/features added.

Some thought will be needed to know what kinds of tests would be appropriate. This mostly affects libraries that use Numba. Maybe a little class that implements __array_ufunc__?

If you are closing an issue please use one of the automatic closing words as noted here: https://help.github.com/articles/closing-issues-using-keywords/

It could be related to #6678, #4625, and/or #2979. For completeness, the PR in its current state only addresses ufunc's __call__ method; perhaps we want accumulate, at, outer, reduce, and reduceat as well. Also, it only addresses np.vectorize; maybe something similar is needed for np.guvectorize.

I'm leaving it in this simple state so that I can get feedback about whether this is the right direction before proceeding to generalize it. (Also, which generalizations are important?)

…all__.

sklam · 2023-06-06T14:56:15Z

@jpivarski, Thank you for the effort. This PR is definitely in the right direction.

Some thought will be needed to know what kinds of tests would be appropriate. This mostly affects libraries that use Numba. Maybe a little class that implements array_ufunc?

Yes, that would do. A minimal class to exercise the new code path in Dufunc would be sufficient.

It could be related to #6678, #4625, and/or #2979. For completeness, the PR in its current state only addresses ufunc's call method; perhaps we want accumulate, at, outer, reduce, and reduceat as well. Also, it only addresses np.vectorize; maybe something similar is needed for np.guvectorize.

I think we should keep this PR small and focus on __array_ufunc__ __call__ support and defer the support for other ufunc methods and guvectorize to different PRs.

In summary, I think all that's remaining for this PR are tests and a doc update.

jpivarski · 2023-06-07T21:53:25Z

I added a test (we'll see what happens in CI), though I'm not sure what to say where for a doc update. I never saw any documentation that says that ufuncs made by @vectorize aren't recognized by NEP13-overloaded arrays; it's something I just ran into. Do you have a suggestion of what text should be added?

Thanks!

sklam · 2023-06-09T14:07:57Z

For docs, I think it's worth adding a note in https://github.com/numba/numba/blob/main/docs/source/reference/jit-compilation.rst#vectorized-functions-ufuncs-and-dufuncs to say what level of NEP13 is supported.

numba/tests/npyufunc/test_ufuncbuilding.py

jpivarski · 2023-06-09T14:57:46Z

Okay, I think that's it: the assertions are unittest-style and I've added a paragraph at the end of the jit-compilation.rst reference.

jpivarski · 2023-07-14T21:52:00Z

Is this waiting on anything from me? As far as I'm aware, it's done.

sklam

Thanks for the contribution!

Allow libraries that implement __array_ufunc__ to override DUFunc.__c…

b8b8365

…all__.

esc added the 2 - In Progress label Jun 2, 2023

Add a test.

5254e9d

jpivarski marked this pull request as ready for review June 7, 2023 21:50

jpivarski requested review from sklam and stuartarchibald as code owners June 7, 2023 21:50

sklam reviewed Jun 9, 2023

View reviewed changes

numba/tests/npyufunc/test_ufuncbuilding.py Outdated Show resolved Hide resolved

jpivarski added 2 commits June 9, 2023 09:25

Use unittest's TestCase.assert*, rather than assert.

8951cd4

Added documentation.

87a7078

sklam added 4 - Waiting on reviewer Waiting for reviewer to respond to author and removed 2 - In Progress labels Jul 19, 2023

sklam approved these changes Jul 19, 2023

View reviewed changes

sklam added 5 - Ready to merge Review and testing done, is ready to merge and removed 4 - Waiting on reviewer Waiting for reviewer to respond to author labels Jul 19, 2023

sklam added this to the Numba 0.58 RC milestone Jul 19, 2023

sklam merged commit fc0371c into numba:main Jul 24, 2023
22 checks passed

gmarkall mentioned this pull request Aug 1, 2023

Fixups for PR#9100 #9112

Merged

jpivarski mentioned this pull request Dec 8, 2023

Support nb.vectorize without type annotations scikit-hep/awkward#1133

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Allow libraries that implement __array_ufunc to override DUFunc.c… #8995

Allow libraries that implement __array_ufunc to override DUFunc.c… #8995

jpivarski commented May 31, 2023

sklam commented Jun 6, 2023

jpivarski commented Jun 7, 2023

sklam commented Jun 9, 2023

jpivarski commented Jun 9, 2023

jpivarski commented Jul 14, 2023

sklam left a comment

Allow libraries that implement __array_ufunc__ to override DUFunc.__c… #8995

Allow libraries that implement __array_ufunc__ to override DUFunc.__c… #8995

Conversation

jpivarski commented May 31, 2023

sklam commented Jun 6, 2023

jpivarski commented Jun 7, 2023

sklam commented Jun 9, 2023

jpivarski commented Jun 9, 2023

jpivarski commented Jul 14, 2023

sklam left a comment

Choose a reason for hiding this comment

Allow libraries that implement __array_ufunc to override DUFunc.c… #8995

Allow libraries that implement __array_ufunc to override DUFunc.c… #8995