diff --git a/pytensor/sparse/math.py b/pytensor/sparse/math.py index 61a9cf0ced..a8ee6e07e6 100644 --- a/pytensor/sparse/math.py +++ b/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]] diff --git a/tests/sparse/test_math.py b/tests/sparse/test_math.py index b06c3247c3..4b43d1aefe 100644 --- a/tests/sparse/test_math.py +++ b/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]])