Give implementation for pointwise

Author: PierreQuinton

src/torchjd/autogram/diagonal_sparse_tensor.py

@@ -2,8 +2,70 @@
 
 import torch
 from torch import Tensor
+from torch.ops import aten
 from torch.utils._pytree import tree_map
 
+# pointwise functions applied to one Tensor with `0.0 → 0`
+_pointwise_functions = {
+    aten.abs.default,
+    aten.abs_.default,
+    aten.absolute.default,
+    aten.absolute_.default,
+    aten.neg.default,
+    aten.neg_.default,
+    aten.negative.default,
+    aten.negative_.default,
+    aten.sign.default,
+    aten.sign_.default,
+    aten.sgn.default,
+    aten.sgn_.default,
+    aten.square.default,
+    aten.square_.default,
+    aten.fix.default,
+    aten.fix_.default,
+    aten.floor.default,
+    aten.floor_.default,
+    aten.ceil.default,
+    aten.ceil_.default,
+    aten.trunc.default,
+    aten.trunc_.default,
+    aten.round.default,
+    aten.round_.default,
+    aten.positive.default,
+    aten.expm1.default,
+    aten.expm1_.default,
+    aten.log1p.default,
+    aten.log1p_.default,
+    aten.sqrt.default,
+    aten.sqrt_.default,
+    aten.sin.default,
+    aten.sin_.default,
+    aten.tan.default,
+    aten.tan_.default,
+    aten.sinh.default,
+    aten.sinh_.default,
+    aten.tanh.default,
+    aten.tanh_.default,
+    aten.asin.default,
+    aten.asin_.default,
+    aten.atan.default,
+    aten.atan_.default,
+    aten.asinh.default,
+    aten.asinh_.default,
+    aten.atanh.default,
+    aten.atanh_.default,
+    aten.erf.default,
+    aten.erf_.default,
+    aten.erfinv.default,
+    aten.erfinv_.default,
+    aten.relu.default,
+    aten.relu_.default,
+    aten.hardtanh.default,
+    aten.hardtanh_.default,
+    aten.leaky_relu.default,
+    aten.leaky_relu_.default,
+}
+
 
 class DiagonalSparseTensor(torch.Tensor):
 
@@ -50,10 +112,19 @@ def to_dense(self) -> Tensor:
         return output
 
     @classmethod
-    def __torch_dispatch__(
-        cls, func: {__name__}, types: Any, args: tuple[()] | Any = (), kwargs: Any = None
-    ):
-        kwargs = kwargs if kwargs else {}
+    def __torch_dispatch__(cls, func: {__name__}, types: Any, args: tuple = (), kwargs: Any = None):
+        kwargs = {} if kwargs is None else kwargs
+
+        # If `func` is a pointwise operator that applies to a single Tensor and such that func(0)=0
+        # Then we can apply the transformation to self._data and wrap the result.
+        if func in _pointwise_functions:
+            assert (
+                isinstance(args, tuple) and len(args) == 1 and func(torch.zeros([])).item() == 0.0
+            )
+            sparse_tensor = args[0]
+            assert isinstance(sparse_tensor, DiagonalSparseTensor)
+            new_data = func(sparse_tensor._data)
+            return DiagonalSparseTensor(new_data, sparse_tensor._v_to_p)
 
         # TODO: Handle batched operations (apply to self._data and wrap)
         # TODO: Handle all operations that can be represented with an einsum by translating them

tests/unit/autogram/test_diagonal_sparse_tensor.py

@@ -1,8 +1,9 @@
 import torch
 from pytest import mark
 from torch.testing import assert_close
+from utils.tensors import randn_, zeros_
 
-from torchjd.autogram.diagonal_sparse_tensor import DiagonalSparseTensor
+from torchjd.autogram.diagonal_sparse_tensor import DiagonalSparseTensor, _pointwise_functions
 
 
 @mark.parametrize(
@@ -18,17 +19,42 @@
     ],
 )
 def test_diagonal_spase_tensor_scalar(shape: list[int]):
-    a = torch.randn(shape)
+    a = randn_(shape)
     b = DiagonalSparseTensor(a, list(range(len(shape))))
 
     assert_close(a, b)
 
 
 @mark.parametrize("dim", [1, 2, 3, 4, 5, 10])
 def test_diag_equivalence(dim: int):
-    a = torch.randn([dim])
+    a = randn_([dim])
     b = DiagonalSparseTensor(a, [0, 0])
 
     diag_a = torch.diag(a)
 
     assert_close(b, diag_a)
+
+
+def test_three_virtual_single_physical():
+    dim = 10
+    a = randn_([dim])
+    b = DiagonalSparseTensor(a, [0, 0, 0])
+
+    expected = zeros_([dim, dim, dim])
+    for i in range(dim):
+        expected[i, i, i] = a[i]
+
+    assert_close(b, expected)
+
+
+@mark.parametrize("func", _pointwise_functions)
+def test_pointwise(func):
+    dim = 100
+    a = randn_([dim])
+    b = DiagonalSparseTensor(a, [0, 0])
+    c = b.to_dense()
+    d = func(b)
+    assert isinstance(d, DiagonalSparseTensor)
+
+    # need to be careful about nans
+    assert_close(d, func(c))