improve

Author: PierreQuinton

src/torchjd/autogram/diagonal_sparse_tensor.py

@@ -125,6 +125,7 @@ def to_dense(self) -> Tensor:
                     identity_matrices[j] = torch.eye(self._v_shape[i], device=device, dtype=dtype)
                 einsum_args += [identity_matrices[j], [first_indices[j], i]]
 
+        # Need to be careful about nans, we would want to get identity times nan.
         output = torch.einsum(*einsum_args, output_indices)
         return output
 

tests/unit/autogram/test_diagonal_sparse_tensor.py

@@ -56,12 +56,11 @@ def test_pointwise(func):
     res = func(b)
     assert isinstance(res, DiagonalSparseTensor)
 
-    # need to be careful about nans
-    assert_close(res, func(c))
+    assert_close(res, func(c), equal_nan=True)
 
 
 @mark.parametrize("func", [torch.mean, torch.sum])
-def test_mean(func):
+def test_unary(func):
     dim = 10
     a = randn_([dim])
     b = DiagonalSparseTensor(a, [0, 0])