Fix SparseDenseMultiply gradient when the dense input is a scalar

Author: Sanjays2402

pytensor/sparse/math.py

@@ -786,7 +786,14 @@ def pullback(self, inputs, outputs, gout):
         (gz,) = gout
         assert psb._is_sparse_variable(x) and psb._is_dense_variable(y)
         assert psb._is_sparse_variable(gz)
-        return y * gz, psb.dense_from_sparse(x * gz)
+        if y.type.ndim == 0:
+            # y was broadcast against every entry of x, so its gradient is the
+            # sum of the contributions of all entries. Reduce on the sparse
+            # product directly instead of densifying it first.
+            gy = sp_sum(x * gz, sparse_grad=True)
+        else:
+            gy = psb.dense_from_sparse(x * gz)
+        return y * gz, gy
 
     def infer_shape(self, node, shapes):
         return [shapes[0]]

tests/sparse/test_math.py

@@ -100,6 +100,20 @@ def test_MulDS(self):
             np.array([[1.0, 2], [3, 0], [0, 6]]),
         )
 
+    @pytest.mark.parametrize("format", ["csc", "csr"])
+    def test_MulSD_scalar_grad(self, format):
+        # Multiplying a sparse matrix by a scalar broadcasts the scalar over
+        # every entry, so the gradient wrt the scalar must be a scalar too.
+        array = np.array([[1.0, 0], [3, 0], [0, 6]])
+        x = as_sparse_variable(as_sparse_format(array, format))
+        y = scalar("y", dtype="float64")
+
+        gy = pytensor.grad(psm.sp_sum(multiply(x, y)), y)
+        assert gy.type.ndim == 0
+
+        f = pytensor.function([y], gy)
+        utt.assert_allclose(array.sum(), f(2.0))
+
     def _testSS(self, op, array1=None, array2=None):
         if array1 is None:
             array1 = np.array([[1.0, 0], [3, 0], [0, 6]])