Add Numba implementation of Blockwise

[ricardoV94]

It uses the machinery developed for RVs and Elemwise. The hard part has to do with multiple number of inputs and numba fussiness.

It also improves Blockwise shape inference based on the infer_shape of the core ops

The small cholesky benchmark I added here test runs 10x faster after this PR on my local machine.

Related Issue

• closes Vectorize follow-up #430

[codecov]

Codecov Report

Attention: Patch coverage is 94.18605% with 5 lines in your changes missing coverage. Please review.

Project coverage is 82.10%. Comparing base (ef97287) to head (661408a).
Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/numba/dispatch/blockwise.py	90.00%	4 Missing ⚠️
pytensor/link/numba/dispatch/random.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1015      +/-   ##
==========================================
+ Coverage   82.09%   82.10%   +0.01%     
==========================================
  Files         183      185       +2     
  Lines       48010    48089      +79     
  Branches     8653     8659       +6     
==========================================
+ Hits        39412    39485      +73     
- Misses       6435     6439       +4     
- Partials     2163     2165       +2     

Files with missing lines	Coverage Δ
pytensor/link/numba/dispatch/__init__.py	100.00% <100.00%> (ø)
pytensor/tensor/blockwise.py	85.50% <100.00%> (+1.29%)	⬆️
pytensor/tensor/random/rewriting/numba.py	100.00% <ø> (ø)
pytensor/tensor/rewriting/numba.py	100.00% <100.00%> (ø)
pytensor/link/numba/dispatch/random.py	58.97% <0.00%> (ø)
pytensor/link/numba/dispatch/blockwise.py	90.00% <90.00%> (ø)

... and 1 file with indirect coverage changes

[ricardoV94]

Getting a failure in the CI that I don't reproduce locally: FAILED tests/link/numba/test_blockwise.py::test_blockwise[Det-True] - TypeError: cannot pickle 'PyCapsule' object

Any idea @jessegrabowski @aseyboldt @lucianopaz (sorry for tag spam)

[aseyboldt]

Maybe locally you end up using a different multiprocessing context? It could be that the pytests run in fork mode and locally it runs in forkserver for some reason?

[ricardoV94]

May be also a numba/scipy version difference

[ricardoV94]

Oh the failure is on the python3.10 CI, it passes on python3.12. numba and scipy versions are equivalent.

[ricardoV94]

The problems seems to be the to_tuple dynamically generated function

[jessegrabowski]

looks great