From 42bc282542bc105565fe1cb519d9941141fce653 Mon Sep 17 00:00:00 2001 From: Sean Law Date: Sun, 13 Sep 2026 21:48:19 -0400 Subject: [PATCH 1/8] Added initial import of numba_cuda_mlir for colab testing --- stumpy/core.py | 6 +++++- stumpy/gpu_aamp.py | 5 ++++- stumpy/gpu_stump.py | 5 ++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/stumpy/core.py b/stumpy/core.py index bb36fa7ee..1358bda2e 100644 --- a/stumpy/core.py +++ b/stumpy/core.py @@ -9,7 +9,11 @@ import warnings import numpy as np -from numba import cuda, njit, prange +from numba import njit, prange +try: + from numba_cuda_mlir import cuda +except ModuleNotFoundError: + from numba import cuda from scipy import linalg from scipy.ndimage import maximum_filter1d, minimum_filter1d from scipy.spatial.distance import cdist diff --git a/stumpy/gpu_aamp.py b/stumpy/gpu_aamp.py index 1aebee850..11c4dda8c 100644 --- a/stumpy/gpu_aamp.py +++ b/stumpy/gpu_aamp.py @@ -6,7 +6,10 @@ import os import numpy as np -from numba import cuda +try: + from numba_cuda_mlir import cuda +except ModuleNotFoundError: + from numba import cuda from . import config, core from .mparray import mparray diff --git a/stumpy/gpu_stump.py b/stumpy/gpu_stump.py index c52176a17..552ec6405 100644 --- a/stumpy/gpu_stump.py +++ b/stumpy/gpu_stump.py @@ -6,7 +6,10 @@ import os import numpy as np -from numba import cuda +try: + from numba_cuda_mlir import cuda +except ModuleNotFoundError: + from numba import cuda from . import config, core from .gpu_aamp import gpu_aamp From 17f019b49efcb841e5a320bc00c16ebd9fd83cb1 Mon Sep 17 00:00:00 2001 From: Sean Law Date: Mon, 14 Sep 2026 06:57:52 -0400 Subject: [PATCH 2/8] Fixed black formatting --- stumpy/core.py | 1 + stumpy/gpu_aamp.py | 1 + stumpy/gpu_stump.py | 1 + 3 files changed, 3 insertions(+) diff --git a/stumpy/core.py b/stumpy/core.py index 1358bda2e..060b4fdd8 100644 --- a/stumpy/core.py +++ b/stumpy/core.py @@ -10,6 +10,7 @@ import numpy as np from numba import njit, prange + try: from numba_cuda_mlir import cuda except ModuleNotFoundError: diff --git a/stumpy/gpu_aamp.py b/stumpy/gpu_aamp.py index 11c4dda8c..155d6ee85 100644 --- a/stumpy/gpu_aamp.py +++ b/stumpy/gpu_aamp.py @@ -6,6 +6,7 @@ import os import numpy as np + try: from numba_cuda_mlir import cuda except ModuleNotFoundError: diff --git a/stumpy/gpu_stump.py b/stumpy/gpu_stump.py index 552ec6405..9b651928e 100644 --- a/stumpy/gpu_stump.py +++ b/stumpy/gpu_stump.py @@ -6,6 +6,7 @@ import os import numpy as np + try: from numba_cuda_mlir import cuda except ModuleNotFoundError: From c34de8d16225b5327facae35813834004f34225c Mon Sep 17 00:00:00 2001 From: Sean Law Date: Wed, 16 Sep 2026 22:33:29 -0400 Subject: [PATCH 3/8] Added new cuda environments to pyproject.toml --- pyproject.toml | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 0bbc3e91c..ef9bfaa6f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,3 +122,23 @@ docs = "cd docs && ./setup.sh" black = 'black --exclude=".*\.ipynb" --extend-exclude=".venv|.pixi" ./' isort = 'isort --profile black --skip .venv --skip .pixi ./' flake8 = 'flake8 --extend-exclude=.venv/,.pixi/ ./' + +[tool.pixi.feature.numba-cuda] +platforms = ["linux-64", "linux-aarch64"] + +[tool.pixi.feature.numba-cuda-mlir] +platforms = ["linux-64", "linux-aarch64"] + +[tool.pixi.feature.numba-cuda.dependencies] +cuda-nvcc = "*" +numba-cuda = "*" +numpy = "<2.5" + +[tool.pixi.feature.numba-cuda-mlir.dependencies] +cuda-nvcc = "*" +numba-cuda-mlir = "*" + +[tool.pixi.environments] +numba-cuda = ["numba-cuda"] +numba-cuda-mlir = ["numba-cuda-mlir"] + From ba0229729a20eaddc64abeee1bdc2c0513184144 Mon Sep 17 00:00:00 2001 From: Sean Law Date: Wed, 16 Sep 2026 23:14:34 -0400 Subject: [PATCH 4/8] Changed tests to use numba-cuda-mlir --- tests/test_core.py | 5 ++++- tests/test_gpu_aamp.py | 6 +++++- tests/test_gpu_aampdist.py | 5 ++++- tests/test_gpu_ostinato.py | 5 ++++- tests/test_gpu_stimp.py | 5 ++++- tests/test_gpu_stump.py | 5 ++++- tests/test_non_normalized_decorator.py | 5 ++++- tests/test_precision.py | 5 ++++- 8 files changed, 33 insertions(+), 8 deletions(-) diff --git a/tests/test_core.py b/tests/test_core.py index 28589cfa3..b8cdca205 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -9,7 +9,10 @@ import numpy.testing as npt import pandas as pd import pytest -from numba import cuda +try: + from numba-cuda-mlir import cuda +except ModuleNotFoundError: + from numba import cuda from scipy.spatial.distance import cdist from stumpy import config, core, rng diff --git a/tests/test_gpu_aamp.py b/tests/test_gpu_aamp.py index 7adfca221..975aff149 100644 --- a/tests/test_gpu_aamp.py +++ b/tests/test_gpu_aamp.py @@ -3,7 +3,11 @@ import numpy as np import numpy.testing as npt import pandas as pd -from numba import cuda + +try: + from numba-cuda-mlir import cuda +except ModuleNotFoundError: + from numba import cuda from stumpy import config, rng diff --git a/tests/test_gpu_aampdist.py b/tests/test_gpu_aampdist.py index bd3b4838d..b5eae05d3 100644 --- a/tests/test_gpu_aampdist.py +++ b/tests/test_gpu_aampdist.py @@ -2,7 +2,10 @@ import numpy as np import numpy.testing as npt -from numba import cuda +try: + from numba-cuda-mlir import cuda +export ModuleNotFoundError: + from numba import cuda if cuda.is_available(): from stumpy.gpu_aampdist import gpu_aampdist diff --git a/tests/test_gpu_ostinato.py b/tests/test_gpu_ostinato.py index 81d0dbf42..12583160d 100644 --- a/tests/test_gpu_ostinato.py +++ b/tests/test_gpu_ostinato.py @@ -2,7 +2,10 @@ import numpy as np import numpy.testing as npt -from numba import cuda +try: + from numba-cuda-mlir import cuda +except ModuleNotFoundError: + from numba import cuda try: from numba.errors import NumbaPerformanceWarning diff --git a/tests/test_gpu_stimp.py b/tests/test_gpu_stimp.py index d93f77a27..cd6850cda 100644 --- a/tests/test_gpu_stimp.py +++ b/tests/test_gpu_stimp.py @@ -3,7 +3,10 @@ import numpy as np import numpy.testing as npt -from numba import cuda +try: + from numba-cuda-mlir import cuda +except ModuleNotFoundError: + from numba import cuda if cuda.is_available(): from stumpy.gpu_stimp import gpu_stimp diff --git a/tests/test_gpu_stump.py b/tests/test_gpu_stump.py index eec2a0cc5..cec08050e 100644 --- a/tests/test_gpu_stump.py +++ b/tests/test_gpu_stump.py @@ -4,7 +4,10 @@ import numpy as np import numpy.testing as npt import pandas as pd -from numba import cuda +try: + from numba-cuda-mlir import cuda +except ModuleNotFoundError: + from numba import cuda from stumpy import config, rng diff --git a/tests/test_non_normalized_decorator.py b/tests/test_non_normalized_decorator.py index 62687e814..71d0b4c4e 100644 --- a/tests/test_non_normalized_decorator.py +++ b/tests/test_non_normalized_decorator.py @@ -3,7 +3,10 @@ import numpy.testing as npt import tornado.ioloop from dask.distributed import Client, LocalCluster -from numba import cuda +try: + from numba-cuda-mlir import cuda +except ModuleNotFoundError: + from numba import cuda from stumpy import core, rng from stumpy.aamp import aamp diff --git a/tests/test_precision.py b/tests/test_precision.py index 4edb3a93d..f213bed95 100644 --- a/tests/test_precision.py +++ b/tests/test_precision.py @@ -6,7 +6,10 @@ import numpy as np import numpy.testing as npt import pytest -from numba import cuda +try: + from numba-cuda-mlir import cuda +except ModuleNotFoundError: + from numba import cuda from stumpy import config, core, rng, sdp From 89a58b56c47d7261ee5e55af1bff0ab04fc2230a Mon Sep 17 00:00:00 2001 From: Sean Law Date: Wed, 16 Sep 2026 23:26:06 -0400 Subject: [PATCH 5/8] Fixed typo --- tests/test_core.py | 3 ++- tests/test_gpu_aamp.py | 2 +- tests/test_gpu_aamp_ostinato.py | 5 ++++- tests/test_gpu_aamp_stimp.py | 5 ++++- tests/test_gpu_aampdist.py | 4 ++-- tests/test_gpu_mpdist.py | 6 +++++- tests/test_gpu_ostinato.py | 3 ++- tests/test_gpu_stimp.py | 3 ++- tests/test_gpu_stump.py | 3 ++- tests/test_non_normalized_decorator.py | 3 ++- tests/test_precision.py | 3 ++- 11 files changed, 28 insertions(+), 12 deletions(-) diff --git a/tests/test_core.py b/tests/test_core.py index b8cdca205..6f779622b 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -9,8 +9,9 @@ import numpy.testing as npt import pandas as pd import pytest + try: - from numba-cuda-mlir import cuda + from numba_cuda_mlir import cuda except ModuleNotFoundError: from numba import cuda from scipy.spatial.distance import cdist diff --git a/tests/test_gpu_aamp.py b/tests/test_gpu_aamp.py index 975aff149..ab6aa6974 100644 --- a/tests/test_gpu_aamp.py +++ b/tests/test_gpu_aamp.py @@ -5,7 +5,7 @@ import pandas as pd try: - from numba-cuda-mlir import cuda + from numba_cuda_mlir import cuda except ModuleNotFoundError: from numba import cuda diff --git a/tests/test_gpu_aamp_ostinato.py b/tests/test_gpu_aamp_ostinato.py index d45d9982c..73587d111 100644 --- a/tests/test_gpu_aamp_ostinato.py +++ b/tests/test_gpu_aamp_ostinato.py @@ -3,7 +3,10 @@ import mersenne import numpy as np import numpy.testing as npt -from numba import cuda +try: + from numba_cuda_mlir import cuda +except ModuleNotFoundError: + from numba import cuda try: from numba.errors import NumbaPerformanceWarning diff --git a/tests/test_gpu_aamp_stimp.py b/tests/test_gpu_aamp_stimp.py index 2f4fe13d5..3d80487e3 100644 --- a/tests/test_gpu_aamp_stimp.py +++ b/tests/test_gpu_aamp_stimp.py @@ -2,7 +2,10 @@ import numpy as np import numpy.testing as npt -from numba import cuda +try: + from numba_cuda_mlir import cuda +except ModuleNotFoundError: + from numba import cuda if cuda.is_available(): from stumpy.gpu_aamp_stimp import gpu_aamp_stimp diff --git a/tests/test_gpu_aampdist.py b/tests/test_gpu_aampdist.py index b5eae05d3..0d4fea246 100644 --- a/tests/test_gpu_aampdist.py +++ b/tests/test_gpu_aampdist.py @@ -3,8 +3,8 @@ import numpy as np import numpy.testing as npt try: - from numba-cuda-mlir import cuda -export ModuleNotFoundError: + from numba_cuda_mlir import cuda +except ModuleNotFoundError: from numba import cuda if cuda.is_available(): diff --git a/tests/test_gpu_mpdist.py b/tests/test_gpu_mpdist.py index 21de32e90..b8acb1fcf 100644 --- a/tests/test_gpu_mpdist.py +++ b/tests/test_gpu_mpdist.py @@ -3,7 +3,11 @@ import numpy as np import numpy.testing as npt -from numba import cuda + +try: + from numba_cuda_mlir import cuda +except ModuleNotFoundError: + from numba import cuda if cuda.is_available(): from stumpy.gpu_mpdist import gpu_mpdist diff --git a/tests/test_gpu_ostinato.py b/tests/test_gpu_ostinato.py index 12583160d..c98f92b0e 100644 --- a/tests/test_gpu_ostinato.py +++ b/tests/test_gpu_ostinato.py @@ -2,8 +2,9 @@ import numpy as np import numpy.testing as npt + try: - from numba-cuda-mlir import cuda + from numba_cuda_mlir import cuda except ModuleNotFoundError: from numba import cuda diff --git a/tests/test_gpu_stimp.py b/tests/test_gpu_stimp.py index cd6850cda..45d1590ea 100644 --- a/tests/test_gpu_stimp.py +++ b/tests/test_gpu_stimp.py @@ -3,8 +3,9 @@ import numpy as np import numpy.testing as npt + try: - from numba-cuda-mlir import cuda + from numba_cuda_mlir import cuda except ModuleNotFoundError: from numba import cuda diff --git a/tests/test_gpu_stump.py b/tests/test_gpu_stump.py index cec08050e..bbc7c41ca 100644 --- a/tests/test_gpu_stump.py +++ b/tests/test_gpu_stump.py @@ -4,8 +4,9 @@ import numpy as np import numpy.testing as npt import pandas as pd + try: - from numba-cuda-mlir import cuda + from numba_cuda_mlir import cuda except ModuleNotFoundError: from numba import cuda diff --git a/tests/test_non_normalized_decorator.py b/tests/test_non_normalized_decorator.py index 71d0b4c4e..5e7c8f664 100644 --- a/tests/test_non_normalized_decorator.py +++ b/tests/test_non_normalized_decorator.py @@ -3,8 +3,9 @@ import numpy.testing as npt import tornado.ioloop from dask.distributed import Client, LocalCluster + try: - from numba-cuda-mlir import cuda + from numba_cuda_mlir import cuda except ModuleNotFoundError: from numba import cuda diff --git a/tests/test_precision.py b/tests/test_precision.py index f213bed95..8c939cf55 100644 --- a/tests/test_precision.py +++ b/tests/test_precision.py @@ -6,8 +6,9 @@ import numpy as np import numpy.testing as npt import pytest + try: - from numba-cuda-mlir import cuda + from numba_cuda_mlir import cuda except ModuleNotFoundError: from numba import cuda From e1ff369024e526164970d1eeafb49d886874acf6 Mon Sep 17 00:00:00 2001 From: Sean Law Date: Wed, 16 Sep 2026 23:28:37 -0400 Subject: [PATCH 6/8] Minor change --- tests/test_gpu_aamp_ostinato.py | 1 + tests/test_gpu_aamp_stimp.py | 1 + tests/test_gpu_aampdist.py | 1 + 3 files changed, 3 insertions(+) diff --git a/tests/test_gpu_aamp_ostinato.py b/tests/test_gpu_aamp_ostinato.py index 73587d111..b5337cfe3 100644 --- a/tests/test_gpu_aamp_ostinato.py +++ b/tests/test_gpu_aamp_ostinato.py @@ -3,6 +3,7 @@ import mersenne import numpy as np import numpy.testing as npt + try: from numba_cuda_mlir import cuda except ModuleNotFoundError: diff --git a/tests/test_gpu_aamp_stimp.py b/tests/test_gpu_aamp_stimp.py index 3d80487e3..fdfac640c 100644 --- a/tests/test_gpu_aamp_stimp.py +++ b/tests/test_gpu_aamp_stimp.py @@ -2,6 +2,7 @@ import numpy as np import numpy.testing as npt + try: from numba_cuda_mlir import cuda except ModuleNotFoundError: diff --git a/tests/test_gpu_aampdist.py b/tests/test_gpu_aampdist.py index 0d4fea246..b92f8eace 100644 --- a/tests/test_gpu_aampdist.py +++ b/tests/test_gpu_aampdist.py @@ -2,6 +2,7 @@ import numpy as np import numpy.testing as npt + try: from numba_cuda_mlir import cuda except ModuleNotFoundError: From f191992f07e28ab41ef6639befe95efed0bdc35f Mon Sep 17 00:00:00 2001 From: Sean Law Date: Wed, 16 Sep 2026 23:42:36 -0400 Subject: [PATCH 7/8] Updated cuda-version --- pyproject.toml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index ef9bfaa6f..14582d5e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -126,15 +126,23 @@ flake8 = 'flake8 --extend-exclude=.venv/,.pixi/ ./' [tool.pixi.feature.numba-cuda] platforms = ["linux-64", "linux-aarch64"] +[tool.pixi.feature.numba-cuda.system-requirements] +cuda = "13.0" + [tool.pixi.feature.numba-cuda-mlir] platforms = ["linux-64", "linux-aarch64"] +[tool.pixi.feature.numba-cuda-mlir.system-requirements] +cuda = "13.0" + [tool.pixi.feature.numba-cuda.dependencies] +cuda-version = "13.0" cuda-nvcc = "*" numba-cuda = "*" numpy = "<2.5" [tool.pixi.feature.numba-cuda-mlir.dependencies] +cuda-version = "13.0" cuda-nvcc = "*" numba-cuda-mlir = "*" From 60562e36e3b650d27daa41b966e907329dfafad5 Mon Sep 17 00:00:00 2001 From: Sean Law Date: Wed, 16 Sep 2026 23:51:24 -0400 Subject: [PATCH 8/8] Removed cuda-version --- pyproject.toml | 8 -------- 1 file changed, 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 14582d5e1..ef9bfaa6f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -126,23 +126,15 @@ flake8 = 'flake8 --extend-exclude=.venv/,.pixi/ ./' [tool.pixi.feature.numba-cuda] platforms = ["linux-64", "linux-aarch64"] -[tool.pixi.feature.numba-cuda.system-requirements] -cuda = "13.0" - [tool.pixi.feature.numba-cuda-mlir] platforms = ["linux-64", "linux-aarch64"] -[tool.pixi.feature.numba-cuda-mlir.system-requirements] -cuda = "13.0" - [tool.pixi.feature.numba-cuda.dependencies] -cuda-version = "13.0" cuda-nvcc = "*" numba-cuda = "*" numpy = "<2.5" [tool.pixi.feature.numba-cuda-mlir.dependencies] -cuda-version = "13.0" cuda-nvcc = "*" numba-cuda-mlir = "*"