meta-issue: Numba/NumPy interation+collaboration+discussion!

@rgommers

Following on from some recent OOB discussions, this is a meta-issue to track ways in which the Numba and NumPy projects could cooperate over common interaction points. The items presented herein are with the sole purpose of attempting to highlight areas to promote discussion and collaboration. Some of the following is also likely to apply to other packages that have similar requirements to Numba (i.e. type stability). Please feel free to comment and/or add more items.

CC @rgommers @seibert

Type stability challenges in the NumPy APIs.

Background: Numba uses type inference to compute the type of all variables in user code. This is performed by combining the type information from the arguments to the user function with a lot of rules defined in Numba about what the types of the variables returned by operations and functions for given input types will be. Note that type inference is performed at compile time and depends solely on the types of the arguments involved and not on their values.

The following contains specific examples of the kinds of problems a type inferring compiler may run into when attempting to replicate the NumPy API. The NumPy API has obviously grown and evolved over the years, it makes a lot of use of Python's dynamic nature, and predates e.g. Numba by a long way, so it is entirely understandable that these challenges exist. What will be interesting is working out if they can be overcome via e.g. more advanced type systems or type stable API updates.

Functions which change their output type depending on value.

Example::

In [12]: np.linalg.eigvals([[1, 1], [1, 1]])
Out[12]: array([2., 0.])

In [13]: np.linalg.eigvals([[1, 1],[-1, 1]])
Out[13]: array([1.+1.j, 1.-1.j])

this example is at one extreme of the issue in that there's a numerical reason for such a behaviour, however it demonstrates that some functions change their return type (float64 vs. complex128 array type) depending on run time value. Another such example is np.cross, which is discussed here #4128 (comment), where there is a return type specialisation depending on the runtime value of an index of an array shape.

Functions which change their output type depending on values of "flags".

Example::

In [18]: np.unique([1, 2, 1], return_index=False)
Out[18]: array([1, 2])

In [19]: np.unique([1, 2, 1], return_index=True)
Out[19]: (array([1, 2]), array([0, 1]))

The problem here is that unless the return_index flag is provided as a compile time constant value (literally True/False in the user code), it's impossible to determine the return type at compile time.

Functions that return mutable heterogeneous containers (e.g. lists) when their size is bounded.

Example::

In [21]: np.broadcast_arrays(np.ones(2), 3)
Out[21]: [array([1., 1.]), array([3, 3])]

In [22]: np.broadcast_arrays(np.ones(2), 3, 2)
Out[22]: [array([1., 1.]), array([3, 3]), array([2, 2])]

The size of the output container is a predicable function of the number of arguments. In a type system such as Numba's there is no support for heterogeneously typed mutable containers, however there is support for heterogeneously typed immutable containers. Given the above is predicable, could the return type reasonably be a tuple?

Functions implementations that do not match docstrings.

Example::

In [38]: print('\n'.join(np.broadcast_shapes.__doc__.splitlines()[:11]))

    Broadcast the input shapes into a single shape.

    :ref:`Learn more about broadcasting here <basics.broadcasting>`.

    .. versionadded:: 1.20.0

    Parameters
    ----------
    `*args` : tuples of ints, or ints
        The shapes to be broadcast against each other.

In [39]: np.broadcast_shapes([12,], (12,))
Out[39]: (12,)

Whilst not directly type stability... the issue is that Numba contributors end up having to "reverse" what is going on to work out what is accepted as input. It makes the support experience in Numba a little inconsistent depending on what was implemented and worked out at the time. There are also occasions where what is accepted changes between NumPy versions.

Some more general mismatches between NumPy and Numba.

There are a few areas where NumPy and Numba differ where they perhaps should not and it'd be good to discuss how to minimise the existing differences and make future ones smaller!

Naming of arguments to functions is sometimes inconsistent.

Example:

In [48]: np.linalg.matrix_power
Out[48]: <function numpy.linalg.matrix_power(a, n)>

In [49]: np.linalg.matrix_rank
Out[49]: <function numpy.linalg.matrix_rank(M, tol=None, hermitian=False)>

these two functions are in the same module and both take a matrix or stack of matrices as the first argument. However one implementation uses a as the argument name and the other M. This doesn't matter hugely but does present a bit of a general inconsistency. This impacts Numba users now and again when they write code which binds to variable names at the call site, if Numba hasn't used the same names as NumPy then the code won't compile. Numba really ought to fix this in its code base, but it might be something NumPy should consider too?

e.g.

In [52]: from numba import njit

In [53]: @njit
    ...: def foo(x):
    ...:     np.linalg.matrix_rank(M=x)
    ...:

In [54]: try:
    ...:     foo(np.ones((2, 2)))
    ...: except Exception as e:
    ...:     print(e)
    ...:
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function matrix_rank at 0x7ffa1b198430>) found for signature:

 >>> matrix_rank(M=array(float64, 2d, C))

There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'matrix_rank_impl': File: numba/np/linalg.py: Line 2419.
    With argument(s): '(M=array(float64, 2d, C))':
   Rejected as the implementation raised a specific error:
     TypingError: missing a required argument: 'a'

Numba erroneously uses a as its argument name, NumPy uses M.

Mismatch in the use of BLAS/LAPACK precision.

The majority of the NumPy linalg routines all upcast float32 and complex64 types to their 64 bit equivalents and then do the work at this precision before casting the answer back to the input type. This creates problems for Numba in that Numba type specialises the BLAS/LAPACK calls based on the input types, which results in a mismatch compared to NumPy's result. This is also why Numba has better performance NumPy in this area when it probably should not. It also causes a degree of issue in Numba's test suite in that the amount that the results differ is sufficiently large to have to write specialised reconstruction routines to check results (opposed to the more usual np.testing.assert_allclose equivalent).

Inf/NaN handling challenges.

Numba has to reimplement NumPy functionality using a what-Numba-supports subset of what is available in NumPy. As a result translating algorithms can be challenging, particularly with regards to Inf/NaN production as it is not always clear whether such an output was intentional or a side effect of the implementation. This often appears as an issue in the case of supporting a new version of NumPy in which some algorithm has changed and Inf/NaN production has changed as a result. An example might be the NumPy 1.19 upgrade #7019 where np.percentile and np.quantile changed production of inf/nan (and also dropped boolean support!). It might be beneficial to specify the expected Inf/NaN handling behaviour of functions cf. standards such as C99 or POSIX.

Other areas of interaction.

Exposing NumPy's internals as ctypes pointers.

Relatively recently NumPy added BitGenerator/Generator to the np.random module. Part of the implementation involved exposing the underlying C routines, that produce the bit streams and the "next" random number of a given type, via ctypes/CFFI. This is very useful on two fronts in
terms of Numba providing support for the same:

It means that Numba doesn't have to literally copy/replicate the entire C implementation in its code base.
It means that Numba can easily achieve bitwise parity to NumPy at a low level as literally the same functions are being called.

Essentially, due to this, it was possible to write something to get basic support plumbed through in a couple of hours (#4499 (comment)), without this it would have probably taken nearer a week, so thank you NumPy engineers, this was great! It also means that if NumPy needs to change something low level, these details are hidden from Numba!

There are a couple potential future improvements to this approach:

The ctypes functions appear as addresses which end up getting baked into the Numba JIT implementation as long longs which breaks Numba's on-disk caching. This is a problem for Numba to resolve!
The ctypes bound functions are essentially called like C function pointers, which means the performance is always going to be limited by having to stage a call and to use the compiled implementation from the NumPy DSO. Were there bitcode equivalants of these functions produced then Numba/JIT compilers could just consume these and LLVM could optimise across these function boundaries.

FFT support in Numba.

Tracking ticket: #5864. The main issue is that there's no binding exposed for the FFT functions in the way there is for BLAS/LAPACK (via SciPy) and now also BitGenerator. As noted on the ticket, part of the issue is that there's no "consistent" interface for dealing with FFTs. Resolving this in a consistent type stable interface and then exposing this via e.g. ctypes would likely help all JIT applications needing FFT support (and also non-JIT applications!).

Ufuncs.

There are a few interaction points in relation to ufuncs that might be of interest.

The ufunc inner loops could be provided as bitcode such that JIT compilers based on LLVM can request the bitcode implementation for given types and just consume that as part of compilation. This has the advantage of NumPy changes being picked up immediately by JIT compilers and that JIT compiler output becomes more consistent with NumPy. It also means a machine specific inner loop will always be generated by the JIT compiler.
NumPy could supply a hook for adding a JIT compiler to the ufunc machinery. As a result, if there is no specific loop available for given input types (i.e. casting would be required to get a match), the JIT compiler could potentially compile one which could lead to improved performance.
NumPy could abstract the entire ufunc interface such that an alternative implementation could be provided. Mixed with ideas in 1. and 2. this would let e.g. a JIT compiler produce CPU specific loops specialised directly on the types present. This again potentially leading to increased performance.

Release testing

Numba's test suite and build farm.

Numba has a large test suite, a lot of which exercises the NumPy APIs. The tests in the test suite are similar to NumPy's but not always the same. It's quite common for Numba's test suite to identify potential bugs/changes made to NumPy across releases.

Numba also has a large CI infrastructure which tests approximately the latest 3 or 4 NumPy releases against the Numba test suite on a large number of platforms. It has been noted a few times previously that it would be good for Numba to test against upcoming NumPy releases early as a mutually beneficial operation. Numba's test suite could potentially pick up issues in NumPy and vice versa. Further, if there are C-API changes in NumPy (especially with respect to ufuncs) there's a reasonable chance that Numba will pick them up. Numba has an "integration" testing project https://github.com/numba/numba-integration-testing that may help with orchestration of this.

@stuartarchibald

Thanks @stuartarchibald, this is a super useful write-up!

Most of this is quite clear, and it seems feasible to make a number of changes in NumPy that would address some of these points. I'd like to open a discussion on the NumPy mailing list soon about putting something like "improve NumPy-Numba interactions" on the roadmap based on the above.

A few questions:

np.broadcast_arrays example [...] In a type system such as Numba's there is no support for heterogeneously typed mutable containers, however there is support for heterogeneously typed immutable containers.

This one surprised me a little. Is that a fundamental problem with Numba's type system? What do you do currently, not support such functions, not make them work in nopython mode, or just return a tuple anyway?

np.broadcast_shapes example [...] Whilst not directly type stability... the issue is that Numba contributors end up having to "reverse" what is going on to work out what is accepted as input.

Is this different from the general "wherever NumPy expects arrays, it also accepts anything that can be coerced to an array via a np.asarray call (sequences, scalars, generators, objects with a __array__ attribute, etc."?

Naming of arguments to functions is sometimes inconsistent.

That one is probably best fixed by making arguments where the one-letter name is not meaningful positional-only. Just to make sure: that would fix the issue for Numba, right?

Mismatch in the use of BLAS/LAPACK precision. [...] which results in a mismatch compared to NumPy's result. This is also why Numba has better performance NumPy in this area when it probably should not.

I thought there was a significant performance gain also because Numba translates np.linalg calls to their scipy.linalg.cython_blas/lapack equivalents? (float64 is up to 4x faster for small inputs IIRC, due to avoiding Python overhead). That should remain even if you'd do an upcast right?

It's a little questionable by the way that NumPy does this upcasting. It won't be changed, but it basically means that there's not much point in using those lower-precision dtypes. They're used much more in deep learning libraries, which don't do this upcast (the lower precision results are fine, and performance gains are important).

@sklam

re: the heterogenously typed mutable containers. The limitation in Numba is not so much the type system as the compilation model that the type system is designed to serve. Numba's JIT compilation model requires functions that can be statically typed at call time. The dispatch mechanism can select between multiple compiled implementations based on argument types, but only as control passes from the Python interpreter into nopython mode. Once in nopython mode, every value (and every function call) needs a static type derivable from those input argument types.

This creates a problem for a mutable container, because the return type of __getitem__() isn't known at compile time, even for constant indices. Numba can make the types of tuple elements part of the type of the tuple itself (as if it were like a C struct), and that lets us handle things like constant indexing or even unrolling a loop over the tuple elements. For a mutable heterogenous container, like a list, compilation would basically have to stop as soon as the list was indexed, and could only continue once the return type of the getitem was known. (I've had conversations with @sklam about what it would take for Numba (or something like it) to support a partial compilation model that would solve this, but it is a radical change to how things currently work.)

We do offer a custom "typed list" (as well as typed dict) user data structure in Numba, but that is by design a homogenously typed list, which avoids these compilation issues.

@stuartarchibald

Thanks @stuartarchibald, this is a super useful write-up!

@rgommers No problem, thanks for taking a look at this.

Most of this is quite clear, and it seems feasible to make a number of changes in NumPy that would address some of these points. I'd like to open a discussion on the NumPy mailing list soon about putting something like "improve NumPy-Numba interactions" on the roadmap based on the above.

A few questions:

np.broadcast_arrays example [...] In a type system such as Numba's there is no support for heterogeneously typed mutable containers, however there is support for heterogeneously typed immutable containers.
This one surprised me a little. Is that a fundamental problem with Numba's type system?

As @seibert notes above, the issue is essentially around non-constant indexing. Were a tuple returned in this case, Numba has literal_unroll which does loop body versioning based on type (docs: https://numba.readthedocs.io/en/stable/reference/pysupported.html#heterogeneous-tuples), so that could be used to iterate over the tuple if needed. Constant indexing on a tuple works too as the type of the index locations is baked into the tuple type itself. As an aside, there's also some benefits of having a homogeneously typed list e.g. SIMD operations can potentially be issued on a list of floats.

What do you do currently, not support such functions, not make them work in nopython mode, or just return a tuple anyway?

In the specific case above, a restriction was placed on the types of the arguments so as to guarantee a homogeneous type for the returned list.

numba/numba/np/arrayobj.py

Lines 1420 to 1433 in e5bbae7

    
           unified_dtype = None 
        
           dt = None 
        
           for arg in args: 
        
               if isinstance(arg, (types.Array, types.BaseTuple)): 
        
                   dt = arg.dtype 
        
               else: 
        
                   dt = arg 
        
               if unified_dtype is None: 
        
                   unified_dtype = dt 
        
               elif unified_dtype != dt: 
        
                   raise errors.TypingError('Mismatch of argument types. Numba cannot ' 
        
                                            'broadcast arrays with different types. ' 
        
                                            f'Got {args}')

Typically, such problems are discussed and then a solution is worked out that preserves the fundamental type that NumPy produces, but this often comes at the cost of making restrictions in the Numba implementation elsewhere.

np.broadcast_shapes example [...] Whilst not directly type stability... the issue is that Numba contributors end up having to "reverse" what is going on to work out what is accepted as input.

Is this different from the general "wherever NumPy expects arrays, it also accepts anything that can be coerced to an array via a np.asarray call (sequences, scalars, generators, objects with a __array__ attribute, etc."?

Somewhat anecdotally, it is in part this coercion through np.asarray, but that appears to not always have been applied consistently through time (it may be more so now), so there has always been an element of "reversing" what is accepted. I would guess that Numba could now use np.asarray more often in its implementations to get closer to matching what NumPy can do.

Naming of arguments to functions is sometimes inconsistent.

That one is probably best fixed by making arguments where the one-letter name is not meaningful positional-only. Just to make sure: that would fix the issue for Numba, right?

I think that would help, but Numba also needs to start being more consistent on this too. I wonder if some convention for argument naming might help further?

Mismatch in the use of BLAS/LAPACK precision. [...] which results in a mismatch compared to NumPy's result. This is also why Numba has better performance NumPy in this area when it probably should not.

I thought there was a significant performance gain also because Numba translates np.linalg calls to their scipy.linalg.cython_blas/lapack equivalents? (float64 is up to 4x faster for small inputs IIRC, due to avoiding Python overhead). That should remain even if you'd do an upcast right?

I suspect the gain is a combination of things, there's certainly no Python overhead as Numba just wires in the library calls. With arrays large enough to reach the capacity of the hardware, I'd expect Numba to perform better in e.g. the float32 case because it would be calling float32 routines whereas NumPy would be calling float64.

It's a little questionable by the way that NumPy does this upcasting. It won't be changed, but it basically means that there's not much point in using those lower-precision dtypes. They're used much more in deep learning libraries, which don't do this upcast (the lower precision results are fine, and performance gains are important).

I'm not sure why this upcasting occurs, but also anticipate that it cannot be changed! Perhaps an entry in the documentation could provide this information to users (assuming there isn't one already, and such information would be useful)?

Glancing over this (and reading the email), my feeling is that there are the two categories:

A new release of NumPy breaking numba (small or big)
API (usually ancient) where API changes or additions would be helpful

I still think that the first is something that numba should and can hopefully look into. By adding a (basic?) CI run against our nightly wheels. For example, I now saw that sklearn has this set-up to not fail CI, but rather post (and I think update) an issue.
(napari has as similar setup for pip --pre, but the "nightly" wheels would be better).

A test setup could have prevent all of the large breaks in the past year, and I think it would also help with the small ones: You can promptly open an issue at NumPy to ask whether a change was intentional or clarify details about a particular change.
(I still often feel that NumPy changes are way too often taken as "given" when they should not be. An example was when tensorflow imagined a NumPy ABI – rather than small API – break and just accepted it without even opening a NumPy issue to ask if it was intentional and notify of problems.)

The second thing: NumPy can maybe find ways to tag/float issues a bit better when there is explicit downstream need for them (this is not limited to numba, if an API addition would help sklearn that obviously is also important).
Of course API changes and all additions always require some champion and unfortunately, I am not sure that NumPy will have the man-power to pick up most of them. But maybe we can work together to make sure such issues become more visible at the very least. (If I understood Ralf correctly, he semi-agreed to champion improving the NumPy tagging as a start.)

OTOH, many of these seem rather mild breaks (e.g. returning a tuple rather than a list). Those could break users, but we can also decide to go through with it with good reasons.
One point about such changes is that it might be nice to pack them up into a "major" release, since they are possibly breaking, but a warning is not possible. (That is not a thing to worry about for Numba, but it makes it important to clearly tag them in NumPy.)

It doesn't seem to have been mentioned in this issue yet, so it's worth noting that quite a few of these issues intersect with changes from the array API https://numpy.org/devdocs/reference/array_api.html. For example

unique is split into different unique_* functions with consistent return types (the shape is still data-dependent, but that's fundamental to the function).
Every function/operation in the array API has a return dtype that depends only on the input dtypes.
Functions use positional-only arguments (this is something I forgot to mention in that doc), which presumably fixes the argument naming issue.
The broadcast_arrays thing actually is a list in the array API. Perhaps it should be changed to a Tuple there.

A challenge as Ralf mentioned on the NumPy mailing list is that the array API is only a subset of NumPy, so if NumPy were to just focus on array API compatibility it would miss similar issues in other places. For example, eigvals is not currently in the array API, and it's possible it won't be added. If it is added, it will not have the issue you mentioned because the array API functions never use value-dependent return dtypes. So as NumPy makes changes to be compatible with the array API, it should make similar changes everywhere else too, when it makes sense.

As an aside, I'm a little confused by the broadcast_shapes question. Is the issue that it says "tuples" but also accepts lists? That seems like a pretty standard thing that any Python function would do. Perhaps a true mypy compatible signature would more accurately say "Sequence[int]" or something. I would be surprised if NumPy restricted the inputs to only tuples for no good reason. OTOH, I would not really be that surprised if numba did require the input to be tuples, because my general experience with Numba is that it is much more restricting about types.

I'll note that fixing all these issues as much as possible is included on the NumPy 2.0 roadmap. Most things cannot be deprecated, instead they must/will simply be changed in NumPy's main branch after the 1.25.x branch split in May/June'23.

WIP: MAINT: change list-of-array to tuple-of-array returns (Numba com…

Functions that return mutable heterogeneous containers (e.g. lists) when their size is bounded.

numpy/numpy#25570 attempts to address this point. It probably requires some more motivation before it can be accepted, because of the backwards compatibility impact - input from the Numba team there is very welcome.

@rgommers

Functions that return mutable heterogeneous containers (e.g. lists) when their size is bounded.

numpy/numpy#25570 attempts to address this point. It probably requires some more motivation before it can be accepted, because of the backwards compatibility impact - input from the Numba team there is very welcome.

Thanks for addressing this @rgommers, I've commented on the PR noted above with some remarks surrounding the benefits of this change from Numba's perspective. Hope this helps with the decision!

WIP: MAINT: change list-of-array to tuple-of-array returns (Numba com…

ENH: change list-of-array to tuple-of-array returns (Numba compat) (#…

Navigation Menu

Repository navigation

meta-issue: Numba/NumPy interation+collaboration+discussion! #8008

Type stability challenges in the NumPy APIs.

Functions which change their output type depending on value.

Functions which change their output type depending on values of "flags".

Functions that return mutable heterogeneous containers (e.g. lists) when their size is bounded.

Functions implementations that do not match docstrings.

Some more general mismatches between NumPy and Numba.

Naming of arguments to functions is sometimes inconsistent.

Mismatch in the use of BLAS/LAPACK precision.

Inf/NaN handling challenges.

Other areas of interaction.

Exposing NumPy's internals as ctypes pointers.

FFT support in Numba.

Ufuncs.

Release testing

Numba's test suite and build farm.

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	unified_dtype = None
	dt = None
	for arg in args:
	if isinstance(arg, (types.Array, types.BaseTuple)):
	dt = arg.dtype
	else:
	dt = arg

	if unified_dtype is None:
	unified_dtype = dt
	elif unified_dtype != dt:
	raise errors.TypingError('Mismatch of argument types. Numba cannot '
	'broadcast arrays with different types. '
	f'Got {args}')

meta-issue: Numba/NumPy interation+collaboration+discussion! #8008

Description

Type stability challenges in the NumPy APIs.

Functions which change their output type depending on value.

Functions which change their output type depending on values of "flags".

Functions that return mutable heterogeneous containers (e.g. lists) when their size is bounded.

Functions implementations that do not match docstrings.

Some more general mismatches between NumPy and Numba.

Naming of arguments to functions is sometimes inconsistent.

Mismatch in the use of BLAS/LAPACK precision.

Inf/NaN handling challenges.

Other areas of interaction.

Exposing NumPy's internals as ctypes pointers.

FFT support in Numba.

Ufuncs.

Release testing

Numba's test suite and build farm.

Activity

rgommers commented on Apr 28, 2022

seibert commented on Apr 28, 2022

stuartarchibald commented on Apr 28, 2022

seberg commented on Jun 8, 2022

asmeurer commented on Jun 8, 2022

rgommers commented on Apr 10, 2023

rgommers commented on Jan 11, 2024

stuartarchibald commented on Jan 17, 2024

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions