Numba sparse: Implement SparseDenseMultiply

Co-authored-by: Jesse Grabowski <[email protected]>

Author: ricardoV94

pytensor/link/numba/dispatch/sparse/__init__.py

@@ -1 +1 @@
-from pytensor.link.numba.dispatch.sparse import basic, variable
+from pytensor.link.numba.dispatch.sparse import basic, math, variable

pytensor/link/numba/dispatch/sparse/math.py

@@ -0,0 +1,90 @@
+from pytensor.link.numba.dispatch import basic as numba_basic
+from pytensor.link.numba.dispatch.basic import register_funcify_default_op_cache_key
+from pytensor.sparse import SparseDenseMultiply, SparseDenseVectorMultiply
+
+
+@register_funcify_default_op_cache_key(SparseDenseMultiply)
+@register_funcify_default_op_cache_key(SparseDenseVectorMultiply)
+def numba_funcify_SparseDenseMultiply(op, node, **kwargs):
+    x, y = node.inputs
+    [z] = node.outputs
+    out_dtype = z.type.dtype
+    format = z.type.format
+    same_dtype = x.type.dtype == out_dtype
+
+    if y.ndim == 0:
+
+        @numba_basic.numba_njit
+        def sparse_multiply_scalar(x, y):
+            if same_dtype:
+                z = x.copy()
+            else:
+                z = x.astype(out_dtype)
+            # Numba doesn't know how to handle in-place mutation / assignment of fields
+            # z.data *= y
+            z_data = z.data
+            z_data *= y
+            return z
+
+        return sparse_multiply_scalar
+
+    elif y.ndim == 1:
+
+        @numba_basic.numba_njit
+        def sparse_dense_multiply(x, y):
+            assert x.shape[1] == y.shape[0]
+            if same_dtype:
+                z = x.copy()
+            else:
+                z = x.astype(out_dtype)
+
+            M, N = x.shape
+            indices = x.indices
+            indptr = x.indptr
+            z_data = z.data
+            if format == "csc":
+                for j in range(0, N):
+                    for i_idx in range(indptr[j], indptr[j + 1]):
+                        z_data[i_idx] *= y[j]
+                return z
+
+            else:
+                for i in range(0, M):
+                    for j_idx in range(indptr[i], indptr[i + 1]):
+                        j = indices[j_idx]
+                        z_data[j_idx] *= y[j]
+
+            return z
+
+        return sparse_dense_multiply
+
+    else:  # y.ndim == 2
+
+        @numba_basic.numba_njit
+        def sparse_dense_multiply(x, y):
+            assert x.shape == y.shape
+            if same_dtype:
+                z = x.copy()
+            else:
+                z = x.astype(out_dtype)
+
+            M, N = x.shape
+            indices = x.indices
+            indptr = x.indptr
+            z_data = z.data
+            if format == "csc":
+                for j in range(0, N):
+                    for i_idx in range(indptr[j], indptr[j + 1]):
+                        i = indices[i_idx]
+                        z_data[i_idx] *= y[i, j]
+                return z
+
+            else:
+                for i in range(0, M):
+                    for j_idx in range(indptr[i], indptr[i + 1]):
+                        j = indices[j_idx]
+                        z_data[j_idx] *= y[i, j]
+
+            return z
+
+        return sparse_dense_multiply

tests/link/numba/sparse/test_basic.py

@@ -194,16 +194,12 @@ def test_sparse_constant(format, cache):
 
     y_test = np.array([np.pi, np.e, np.euler_gamma])
     with config.change_flags(numba__cache=cache):
-        with pytest.warns(
-            UserWarning,
-            match=r"Numba will use object mode to run SparseDenseVectorMultiply's perform method",
-        ):
-            compare_numba_and_py_sparse(
-                [y],
-                [out],
-                [y_test],
-                eval_obj_mode=False,
-            )
+        compare_numba_and_py_sparse(
+            [y],
+            [out],
+            [y_test],
+            eval_obj_mode=False,
+        )
 
 
 @pytest.mark.parametrize("format", ["csc", "csr"])
@@ -248,15 +244,11 @@ def test_simple_graph(format):
     x_test = sp.sparse.random(3, 3, density=0.5, format=format, random_state=rng)
     y_test = rng.normal(size=(3,))
 
-    with pytest.warns(
-        UserWarning,
-        match=r"Numba will use object mode to run SparseDenseVectorMultiply's perform method",
-    ):
-        compare_numba_and_py_sparse(
-            [x, y],
-            z,
-            [x_test, y_test],
-        )
+    compare_numba_and_py_sparse(
+        [x, y],
+        z,
+        [x_test, y_test],
+    )
 
 
 @pytest.mark.parametrize("format", ("csr", "csc"))

tests/link/numba/sparse/test_math.py

@@ -0,0 +1,28 @@
+import numpy as np
+import pytest
+import scipy
+
+import pytensor.sparse as ps
+import pytensor.tensor as pt
+from tests.link.numba.sparse.test_basic import compare_numba_and_py_sparse
+
+
+pytestmark = pytest.mark.filterwarnings("error")
+
+
+@pytest.mark.parametrize("format", ["csr", "csc"])
+@pytest.mark.parametrize("y_ndim", [0, 1, 2])
+def test_sparse_dense_multiply(y_ndim, format):
+    x = ps.matrix(format, name="x", shape=(3, 3))
+    y = pt.tensor("y", shape=(3,) * y_ndim)
+    z = x * y
+
+    rng = np.random.default_rng((155, y_ndim, format == "csr"))
+    x_test = scipy.sparse.random(3, 3, density=0.5, format=format, random_state=rng)
+    y_test = rng.normal(size=(3,) * y_ndim)
+
+    compare_numba_and_py_sparse(
+        [x, y],
+        z,
+        [x_test, y_test],
+    )