Use objmode in scipy.special without numba-scipy

[aseyboldt]

Here are a few important guidelines and requirements to check before your PR can be merged:

• There is an informative high-level description of the changes.
• The description and/or commit message(s) references the relevant GitHub issue(s).
• pre-commit is installed and set up.
• The commit messages follow these guidelines.
• The commits correspond to relevant logical changes, and there are no commits that fix changes introduced by other commits in the same branch/BR.
• There are tests covering the changes introduced in the PR.

Replaces the numba-scipy related part in #1062.

This adds a new config option numba_scipy that indicates if the numba backend uses numba-scipy for scipy.special functions.
If it is not enabled or if numba_scipy can't be imported it will fall back to objectmode. This makes sure we do not end up with hard to understand numba typing errors if a user didn't install numba-scipy, but instead just a warning about objectmode and somewhat bad performance. Previously these functions would work in the test environment, because numba_scipy is installed there, but not on a normal user install, because numba_scipy is not a dependency of aesara.

I added a new function generate_fallback_impl, extracted from the fallback numba_functify, so that we can call this from within specializations.

I also changed numba_funcify_Elemwise a bit. Previously we would pass the node of the elemwise op to the numba_funcify calls of the scalar_op. Now we create a new node if the op isn't a Composite (we don't need a op there, because Composite.fgraph contains all we need).

[brandonwillard]

What is the cython_support machinery for? At first glance, it looks like some low-level things that were previously handled by numba-scipy are being brought into this project. We definitely don't want to bring that functionality here; it's better that we address such things in numba-scipy itself, whenever possible.

[aseyboldt]

@brandonwillard I had the same thought at first, but I'm not really sure how this would fit anywhere into either numba_scipy or numba itself.

The main complication we have, is that we don't just want exact matches between the special functions and their types, but we also want to allow some casting. So if we ask for a single precision function, we are also ok with a double precision implementation if that's the only thing that's available. A bit (but probably not too much) of the code here could probably still live in numba_scipy if we wanted, but then we'd still have to import it. I don't really like that either however:

In general I'm not really in favor of the general approach of numba_scipy, to just register numba implementations for scipy functions, because this modifies global state. Just importing numba_scipy in a library can completely change the behavior of a different library, because suddenly the typing behavior of numba changes globally.

[codecov]

Codecov Report

Merging #1078 (ceb5741) into main (eace7f6) will increase coverage by 5.16%.
The diff coverage is 85.32%.

❗ Current head ceb5741 differs from pull request most recent head e8eed6c. Consider uploading reports for the commit e8eed6c to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1078      +/-   ##
==========================================
+ Coverage   74.12%   79.29%   +5.16%     
==========================================
  Files         174      153      -21     
  Lines       48654    48099     -555     
  Branches    10353    10937     +584     
==========================================
+ Hits        36064    38138    +2074     
+ Misses      10300     7452    -2848     
- Partials     2290     2509     +219     

Impacted Files	Coverage Δ
aesara/link/numba/dispatch/scalar.py	93.16% <75.00%> (+5.75%)	⬆️
aesara/link/numba/dispatch/cython_support.py	86.95% <86.95%> (ø)
aesara/link/numba/dispatch/basic.py	92.10% <100.00%> (-0.58%)	⬇️
aesara/link/numba/dispatch/elemwise.py	97.16% <100.00%> (+0.04%)	⬆️
aesara/tensor/basic_opt.py	85.90% <0.00%> (-14.10%)	⬇️
aesara/tensor/subtensor_opt.py	87.12% <0.00%> (-12.88%)	⬇️
aesara/tensor/math_opt.py	87.27% <0.00%> (-12.73%)	⬇️
aesara/sparse/type.py	74.76% <0.00%> (-9.64%)	⬇️
aesara/tensor/opt_uncanonicalize.py	96.03% <0.00%> (-3.97%)	⬇️
aesara/link/vm.py	90.26% <0.00%> (-2.30%)	⬇️
... and 106 more

[brandonwillard]

@brandonwillard I had the same thought at first, but I'm not really sure how this would fit anywhere into either numba_scipy or numba itself.

The main complication we have, is that we don't just want exact matches between the special functions and their types, but we also want to allow some casting. So if we ask for a single precision function, we are also ok with a double precision implementation if that's the only thing that's available. A bit (but probably not too much) of the code here could probably still live in numba_scipy if we wanted, but then we'd still have to import it. I don't really like that either however:

In general I'm not really in favor of the general approach of numba_scipy, to just register numba implementations for scipy functions, because this modifies global state. Just importing numba_scipy in a library can completely change the behavior of a different library, because suddenly the typing behavior of numba changes globally.

The point of numba-scipy is to provide implementations for scipy functions, so all these concerns definitely fall within numba-scipy's purview; however, there's a lot missing from it, so we really need to think in terms of how it should be working, and make the relevant changes in numba-scipy. This is especially important due to numba-scipy being more visible and applicable to others in need of scipy support within Numba. If we do any significant amount of work in this direction within Aesara, the results are only applicable to Aesara, and we're nowhere near as likely to get contributions from outside.

From what I can tell, the code in cython_support looks very similar in nature to the code in numba-scipy, and, if it's an improvement over numba-scipy, it shouldn't be restricted to Aesara. More importantly, a PR to numba-scipy will get the attention of others who know Numba well enough to provide useful input about this approach.

[aseyboldt]

I'll try if I can figure out how to put the improved cython support into a PR to numba itself, maybe that actually is the correct place for this. If this works out we might actually still push something to aesara in the meantime, numba just had a release and they don't seem to make releases all that often.

[brandonwillard]

I'll try if I can figure out how to put the improved cython support into a PR to numba itself, maybe that actually is the correct place for this. If this works out we might actually still push something to aesara in the meantime, numba just had a release and they don't seem to make releases all that often.

numba-scipy is a much smaller project, and I believe it can cut releases more easily than Numba.