Test Numba support across the entire test suite

.github/workflows/test.yml

@@ -66,7 +66,7 @@ jobs:
     runs-on: ubuntu-latest
     if: ${{ needs.changes.outputs.changes == 'true' && needs.style.result == 'success' }}
     strategy:
-      fail-fast: true
+      fail-fast: false
       matrix:
         python-version: ["3.7", "3.9"]
         fast-compile: [0]
@@ -115,10 +115,10 @@ jobs:
       - name: Install dependencies
         shell: bash -l {0}
         run: |
-          mamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov sympy
-          if [[ $INSTALL_NUMBA == "1" ]]; then mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.55" numba-scipy; fi
-          mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib
-          pip install -e ./
+          mamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl "numpy>=1.23.3" scipy pip mkl-service graphviz cython pytest coverage pytest-cov sympy filelock etuples logical-unification miniKanren cons typing_extensions "setuptools>=48.0.0"
+          if [[ $INSTALL_NUMBA == "1" ]]; then mamba install --yes -q -c numba/label/dev "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57.dev0"; fi
+          mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numpy>=1.23.3" jax jaxlib
+          pip install --no-deps -e ./
           mamba list && pip freeze
           python -c 'import aesara; print(aesara.config.__str__(print_doc=False))'
           python -c 'import aesara; assert(aesara.config.blas__ldflags != "")'
@@ -132,7 +132,7 @@ jobs:
           if [[ $FAST_COMPILE == "1" ]]; then export AESARA_FLAGS=$AESARA_FLAGS,mode=FAST_COMPILE; fi
           if [[ $FLOAT32 == "1" ]]; then export AESARA_FLAGS=$AESARA_FLAGS,floatX=float32; fi
           export AESARA_FLAGS=$AESARA_FLAGS,warn__ignore_bug_before=all,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe
-          python -m pytest -x -r A --verbose --runslow --cov=aesara/ --cov-report=xml:coverage/coverage-${MATRIX_ID}.xml --no-cov-on-fail $PART
+          python -m pytest --verbose --runslow --cov=aesara/ --cov-report=xml:coverage/coverage-${MATRIX_ID}.xml --no-cov-on-fail $PART
         env:
           MATRIX_ID: ${{ steps.matrix-id.outputs.id }}
           MKL_THREADING_LAYER: GNU

aesara/compile/mode.py

@@ -42,9 +42,9 @@
     "c": CLinker(),  # Don't support gc. so don't check allow_gc
     "c|py": OpWiseCLinker(),  # Use allow_gc Aesara flag
     "c|py_nogc": OpWiseCLinker(allow_gc=False),
-    "vm": VMLinker(use_cloop=False),  # Use allow_gc Aesara flag
+    "vm": NumbaLinker(),  # VMLinker(use_cloop=False),  # Use allow_gc Aesara flag
     "cvm": VMLinker(use_cloop=True),  # Use allow_gc Aesara flag
-    "vm_nogc": VMLinker(allow_gc=False, use_cloop=False),
+    "vm_nogc": NumbaLinker(),  # VMLinker(allow_gc=False, use_cloop=False),
     "cvm_nogc": VMLinker(allow_gc=False, use_cloop=True),
     "jax": JAXLinker(),
     "numba": NumbaLinker(),
@@ -441,9 +441,9 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
 # Use VM_linker to allow lazy evaluation by default.
 FAST_COMPILE = Mode(VMLinker(use_cloop=False, c_thunks=False), "fast_compile")
 if config.cxx:
-    FAST_RUN = Mode("cvm", "fast_run")
+    FAST_RUN = Mode("numba", "fast_run")
 else:
-    FAST_RUN = Mode("vm", "fast_run")
+    FAST_RUN = Mode("numba", "fast_run")
 
 JAX = Mode(
     JAXLinker(),

aesara/configdefaults.py

@@ -392,7 +392,7 @@ def add_compile_configvars():
     config.add(
         "mode",
         "Default compilation mode",
-        ConfigParam("Mode", apply=_filter_mode),
+        ConfigParam("NUMBA", apply=_filter_mode),
         in_c_key=False,
     )
 
@@ -463,7 +463,18 @@ def add_compile_configvars():
             "linker",
             "Default linker used if the aesara flags mode is Mode",
             EnumStr(
-                "cvm", ["c|py", "py", "c", "c|py_nogc", "vm", "vm_nogc", "cvm_nogc"]
+                "numba",
+                [
+                    "c|py",
+                    "py",
+                    "c",
+                    "c|py_nogc",
+                    "vm",
+                    "vm_nogc",
+                    "cvm_nogc",
+                    "numba",
+                    "jax",
+                ],
             ),
             in_c_key=False,
         )
@@ -473,7 +484,7 @@ def add_compile_configvars():
         config.add(
             "linker",
             "Default linker used if the aesara flags mode is Mode",
-            EnumStr("vm", ["py", "vm_nogc"]),
+            EnumStr("numba", ["py", "vm_nogc", "vm", "numba", "jax"]),
             in_c_key=False,
         )
         if type(config).cxx.is_default:

aesara/link/numba/dispatch/basic.py

@@ -32,6 +32,7 @@
 from aesara.scalar.math import Softplus
 from aesara.tensor.blas import BatchedDot
 from aesara.tensor.math import Dot
+from aesara.tensor.random.type import RandomGeneratorType
 from aesara.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 from aesara.tensor.slinalg import Cholesky, Solve
 from aesara.tensor.subtensor import (
@@ -93,6 +94,8 @@ def get_numba_type(
         dtype = np.dtype(aesara_type.dtype)
         numba_dtype = numba.from_dtype(dtype)
         return numba_dtype
+    elif isinstance(aesara_type, RandomGeneratorType):
+        return numba.types.npy_rng
     else:
         raise NotImplementedError(f"Numba type not implemented for {aesara_type}")