Add tests for: SparseFromDense, StructuredDotGrad, and SpSum in numba backend

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

@@ -10,11 +10,7 @@ from pytensor.link.numba.dispatch.basic import (
     register_funcify_default_op_cache_key,
 )
 from pytensor.link.numba.dispatch.compile_ops import numba_deepcopy
-from pytensor.link.numba.dispatch.sparse.variable import (
-    CSMatrixType,
-    csc_matrix_from_components,
-    csr_matrix_from_components,
-)
+from pytensor.link.numba.dispatch.sparse.variable import CSMatrixType
 from pytensor.sparse import (
     CSM,
     Cast,

tests/link/numba/sparse/test_basic.py

@@ -299,3 +299,17 @@ def test_sparse_conversion():
                 res = to_csc(x_inp)
             assert res.format == output_format
             np.testing.assert_array_equal(res.todense(), x_dense)
+
+
+@pytest.mark.parametrize("format", ("csr", "csc"))
+def test_sparse_from_dense(format):
+    x = pt.matrix(name="x", shape=(11, 7), dtype=config.floatX)
+    x_sparse = sp.sparse.random(11, 7, density=0.4, format=format, dtype=config.floatX)
+    x_test = x_sparse.toarray()
+
+    if format == "csr":
+        y = ps.csr_from_dense(x)
+    else:
+        y = ps.csc_from_dense(x)
+
+    compare_numba_and_py_sparse([x], y, [x_test])

tests/link/numba/sparse/test_math.py

@@ -95,3 +95,38 @@ def test_sparse_spmv(sp_format):
     y_test = rng.normal(size=(6,))
 
     compare_numba_and_py_sparse([x, y], z, [x_test, y_test])
+
+
+@pytest.mark.parametrize("x_format", ["csr", "csc"])
+@pytest.mark.parametrize("y_format", ["csr", "csc", "dense"])
+@pytest.mark.parametrize("x_shape, y_shape", DOT_SHAPES)
+def test_structured_dot_grad(x_format, y_format, x_shape, y_shape):
+    rng = np.random.default_rng()
+    g_xy_shape = (x_shape[0], y_shape[1])
+
+    x = ps.matrix(format=x_format, name="x", shape=x_shape)
+    x_test = scipy.sparse.random(*x_shape, density=0.4, format=x_format)
+
+    if y_format == "dense":
+        y = pt.matrix("y", shape=y_shape)
+        g_xy = pt.matrix(name="g_xy", shape=g_xy_shape)
+        y_test = rng.normal(size=y_shape)
+        g_xy_test = rng.normal(size=g_xy_shape)
+    else:
+        y = ps.matrix(format=y_format, name="y", shape=y_shape)
+        g_xy = ps.matrix(format=x_format, name="g_xy", shape=g_xy_shape)
+        y_test = scipy.sparse.random(*y_shape, density=0.5, format=y_format)
+        g_xy_test = scipy.sparse.random(*g_xy_shape, density=0.3, format=x_format)
+
+    z = ps.structured_dot_grad(x, y, g_xy)
+    compare_numba_and_py_sparse([x, y, g_xy], z, [x_test, y_test, g_xy_test])
+
+
+@pytest.mark.parametrize("format", ["csr", "csc"])
+@pytest.mark.parametrize("axis", [None, 0, 1])
+def test_sparse_sum(format, axis):
+    x = ps.matrix(format=format, name="x", shape=(7, 5))
+    z = ps.sp_sum(x, axis=axis)
+    x_test = scipy.sparse.random(7, 5, density=0.4, format=format)
+
+    compare_numba_and_py_sparse([x], z, [x_test])