Added initial import of numba_cuda_mlir for colab testing

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"]
+

stumpy/core.py

@@ -9,7 +9,12 @@
 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

stumpy/gpu_aamp.py

@@ -6,7 +6,11 @@
 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

stumpy/gpu_stump.py

@@ -6,7 +6,11 @@
 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

tests/test_core.py

@@ -9,7 +9,11 @@
 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

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
 

tests/test_gpu_aamp_ostinato.py

@@ -3,7 +3,11 @@
 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

tests/test_gpu_aamp_stimp.py

@@ -2,7 +2,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_aamp_stimp import gpu_aamp_stimp

tests/test_gpu_aampdist.py

@@ -2,7 +2,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_aampdist import gpu_aampdist

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

tests/test_gpu_ostinato.py

@@ -2,7 +2,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
 
 try:
     from numba.errors import NumbaPerformanceWarning

tests/test_gpu_stimp.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_stimp import gpu_stimp

tests/test_gpu_stump.py

@@ -4,7 +4,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
 

tests/test_non_normalized_decorator.py

@@ -3,7 +3,11 @@
 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

tests/test_precision.py

@@ -6,7 +6,11 @@
 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