From 8b444888c032ef4387cc2fe172ff7e4df44b4e52 Mon Sep 17 00:00:00 2001 From: "Brandon T. Willard" Date: Wed, 9 Nov 2022 18:24:48 -0600 Subject: [PATCH] Make Numba the default backend for testing Do not merge this commit; it is intended for temporary use in a draft PR. --- .github/workflows/test.yml | 4 ++-- aesara/compile/mode.py | 8 ++++---- aesara/configdefaults.py | 17 ++++++++++++++--- tests/scan/test_basic.py | 20 +++++++++++++++----- 4 files changed, 35 insertions(+), 14 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 77b9d60f66..903bb83a8a 100644 --- a/.github/workflows/test.yml +++ b/.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: true matrix: python-version: ["3.7", "3.9"] fast-compile: [0] @@ -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 diff --git a/aesara/compile/mode.py b/aesara/compile/mode.py index 5232163ba2..0f43d46c07 100644 --- a/aesara/compile/mode.py +++ b/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(), diff --git a/aesara/configdefaults.py b/aesara/configdefaults.py index b1e914b2a9..f7c323e15e 100644 --- a/aesara/configdefaults.py +++ b/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: diff --git a/tests/scan/test_basic.py b/tests/scan/test_basic.py index b4af67d572..a518961d46 100644 --- a/tests/scan/test_basic.py +++ b/tests/scan/test_basic.py @@ -248,7 +248,7 @@ class TestScan: "rng_type", [ np.random.default_rng, - np.random.RandomState, + # np.random.RandomState, ], ) def test_inner_graph_cloning(self, rng_type): @@ -396,7 +396,7 @@ def f_pow2(x_tm1): assert all(i.value is None for i in scan_node.op.fn.input_storage) assert all(o.value is None for o in scan_node.op.fn.output_storage) - @pytest.mark.parametrize("mode", [Mode(linker="py"), Mode(linker="cvm")]) + @pytest.mark.parametrize("mode", ["NUMBA", Mode(linker="py"), Mode(linker="cvm")]) @pytest.mark.parametrize( "x_init", [ @@ -421,7 +421,12 @@ def f_pow(x_tm1): assert res.dtype == exp_res.dtype @pytest.mark.parametrize( - "mode", [Mode(linker="py", optimizer=None), Mode(linker="cvm", optimizer=None)] + "mode", + [ + "NUMBA", + Mode(linker="py", optimizer=None), + Mode(linker="cvm", optimizer=None), + ], ) @pytest.mark.parametrize( "x", @@ -459,7 +464,12 @@ def inner_fn(x_seq, x_i): assert res.dtype == exp_res.dtype @pytest.mark.parametrize( - "mode", [Mode(linker="py", optimizer=None), Mode(linker="cvm", optimizer=None)] + "mode", + [ + "NUMBA", + Mode(linker="py", optimizer=None), + Mode(linker="cvm", optimizer=None), + ], ) @pytest.mark.parametrize( "x", @@ -1126,7 +1136,7 @@ def test_inner_grad(self): utt.assert_allclose(out, vR) @pytest.mark.parametrize( - "mode", [Mode(linker="cvm", optimizer=None), Mode(linker="cvm")] + "mode", ["NUMBA", Mode(linker="cvm", optimizer=None), Mode(linker="cvm")] ) def test_sequence_is_scan(self, mode): """Make sure that a `Scan` can be used as a sequence input to another `Scan`."""