[WIP] Fix jac undefined errors

Author: ValerianRey

src/torchjd/utils/_accumulation.py

@@ -1,20 +1,24 @@
 from collections.abc import Iterable
+from typing import cast
 
 from torch import Tensor
 
+from torchjd.utils._tensor_with_jac import TensorWithJac
+
 
 def _accumulate_jacs(params: Iterable[Tensor], jacobians: Iterable[Tensor]) -> None:
     for param, jac in zip(params, jacobians, strict=True):
         _check_expects_grad(param)
-        if hasattr(param, "jac"):
-            param.jac += jac
+        if hasattr(param, "jac"):  # No check for None because jac cannot be None
+            param_ = cast(TensorWithJac, param)
+            param_.jac += jac
         else:
             # TODO: this could be a serious memory issue
             # We clone the value because we do not want subsequent accumulations to also affect
             # this value (in case it is still used outside). We do not detach from the
             # computation graph because the value can have grad_fn that we want to keep track of
             # (in case it was obtained via create_graph=True and a differentiable aggregator).
-            param.jac = jac.clone()
+            param.__setattr__("jac", jac.clone())
 
 
 def _accumulate_grads(params: Iterable[Tensor], gradients: Iterable[Tensor]) -> None:

src/torchjd/utils/_jac_to_grad.py

@@ -1,10 +1,12 @@
 from collections.abc import Iterable
+from typing import cast
 
 import torch
 from torch import Tensor
 
 from torchjd.aggregation import Aggregator
 from torchjd.utils._accumulation import _accumulate_grads
+from torchjd.utils._tensor_with_jac import TensorWithJac
 
 
 def jac_to_grad(
@@ -20,17 +22,19 @@ def jac_to_grad(
     :param retain_jacs: Whether to preserve the ``.jac`` fields of the parameters.
     """
 
-    params_ = list(params)
+    params_ = list[TensorWithJac]()
+    for p in params:
+        if not hasattr(p, "jac"):
+            raise ValueError(
+                "Some `jac` fields were not populated. Did you use `autojac.backward` before"
+                "calling `jac_to_grad`?"
+            )
+        p_ = cast(TensorWithJac, p)
+        params_.append(p_)
 
     if len(params_) == 0:
         return
 
-    if not all([hasattr(p, "jac") and p.jac is not None for p in params_]):
-        raise ValueError(
-            "Some `jac` fields were not populated. Did you use `autojac.backward` before calling "
-            "`jac_to_grad`?"
-        )
-
     jacobians = [p.jac for p in params_]
 
     # TODO: check that the Jacobian shapes match
@@ -51,7 +55,7 @@ def _unite_jacobians(jacobians: list[Tensor]) -> Tensor:
 
 
 def _disunite_gradient(
-    gradient_vector: Tensor, jacobians: list[Tensor], params: list[Tensor]
+    gradient_vector: Tensor, jacobians: list[Tensor], params: list[TensorWithJac]
 ) -> list[Tensor]:
     gradient_vectors = []
     start = 0
@@ -64,7 +68,7 @@ def _disunite_gradient(
     return gradients
 
 
-def _free_jacs(params: Iterable[Tensor]) -> None:
+def _free_jacs(params: Iterable[TensorWithJac]) -> None:
     """
     Deletes the ``.jac`` field of the provided parameters.
 

src/torchjd/utils/_tensor_with_jac.py

@@ -0,0 +1,11 @@
+from torch import Tensor
+
+
+class TensorWithJac(Tensor):
+    """
+    Tensor known to have a populated jac field.
+
+    Should not be directly instantiated, but can be used as a type hint and can be casted to.
+    """
+
+    jac: Tensor