fix(autojac): Replace unsafe view with reshape (#582)

* Add test for jac_to_grad with non-contiguous jac (that would have failed in v0.8.1)
* Replace another potentially unsafe view with reshape

Author: DJKorchinski

src/torchjd/autojac/_jac_to_grad.py

@@ -92,7 +92,7 @@ def _disunite_gradient(
     tensors: list[TensorWithJac],
 ) -> list[Tensor]:
     gradient_vectors = gradient_vector.split([t.numel() for t in tensors])
-    gradients = [g.view(t.shape) for g, t in zip(gradient_vectors, tensors, strict=True)]
+    gradients = [g.reshape(t.shape) for g, t in zip(gradient_vectors, tensors, strict=True)]
     return gradients
 
 

tests/unit/autojac/test_jac_to_grad.py

@@ -101,3 +101,17 @@ def test_jacs_are_freed(retain_jac: bool):
     check = assert_has_jac if retain_jac else assert_has_no_jac
     check(t1)
     check(t2)
+
+
+def test_noncontiguous_jac():
+    """Tests that jac_to_grad works when the .jac field is non-contiguous."""
+
+    aggregator = UPGrad()
+    t = tensor_([2.0, 3.0, 4.0], requires_grad=True)
+    jac_T = tensor_([[-4.0, 1.0], [1.0, 6.0], [1.0, 1.0]])
+    jac = jac_T.T
+    t.__setattr__("jac", jac)
+    g = aggregator(jac)
+
+    jac_to_grad([t], aggregator)
+    assert_grad_close(t, g)