Add support for random Generators in Numba backend

[ricardoV94]

Description

TODO:

• Merge code duplication into elemwise_codegen
• Handle broadcast via size argument
• Handle more than normal
• Add expand_dims in make_node. This is actually something that ends up being useful in other places, as in vectorization and dim analysis. It also canonicalizes graphs that are equivalent, as the expand_dims is always implict
• Handle non-scalar inputs / outputs. This would also solve Blockwise for Numba backend!
• Find a way to copy RNGs if inplace=False. Not a huge priority, but a nice to have for consistency
• Remove RandomState from codebase. I vote for no deprecation as this is functionality that was broken in Numba anyway (no way to reseed), and non-default for other backends.

Related Issue

• Closes Use Numba Generators for random graphs and deprecate shared RandomState variables #316
• Closes Get rid of dummy type input to RVs #701
• Closes RandomVariable treats size=() as size=None #568

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[ricardoV94]

Size and rng copy are now supported

[ricardoV94]

@aseyboldt Elemwise is working again in case you want to benchmark before/after

[ricardoV94]

Here is a test snippet:

%env NUMBA_DISABLE_ERROR_MESSAGE_HIGHLIGHTING=1
import colorama
import pytensor
import pytensor.tensor as pt
import numpy as np

x = pt.vector("x", shape=(5,), dtype="float64")
out = pt.psi(x) * 2
fn = pytensor.function([x], out, mode="NUMBA")

x_test = np.random.uniform(size=5)
fn(x_test)

Not setting that variable will lead to the warning showing up (if colorama can be imported)

[ricardoV94]

Tests pass with a77df24

My only concern is that users with colorama will see the warning. I don't know how to filter that kind of messages, and also agree that filtering any NumbaWarning is a tad too much. Should we perhaps ask Numba to have a subclass for caching warnings? Perhaps that would be more reasonable to filter by default?

For clarity there's a sort of regression when we introduced the store_core_outputs for Elemwise, which now shows a colorized cache warning. The filtering of the numba_funcified_fgraph cache warning still works fine. I couldn't figure out why does one work but not the other.

[ricardoV94]

Ok fixed with some regex, I was being tripped by multiple bold escape codes...

[ricardoV94]

Marking as a draft because I want to merge a bugfix and do a release before this major changes get in

[aseyboldt]

Sounds like you had a fun afternoon.
But that explains the strange failures....

[twiecki]

Good one!