Merge branch 'main' into stationarity_property

Author: ValerianRey

tests/unit/autojac/_transform/test_interactions.py

@@ -1,4 +1,5 @@
 import torch
+from pytest import raises
 from torch.testing import assert_close
 
 from torchjd.autojac._transform import (
@@ -248,3 +249,17 @@ def test_equivalence_jac_grads():
     assert_close(jac_A, torch.stack([grad_1_A, grad_2_A]))
     assert_close(jac_b, torch.stack([grad_1_b, grad_2_b]))
     assert_close(jac_c, torch.stack([grad_1_c, grad_2_c]))
+
+
+def test_stack_different_required_keys():
+    """Tests that the Stack transform fails on transforms with different required keys."""
+
+    a = torch.tensor(1.0, requires_grad=True)
+    y1 = a * 2.0
+    y2 = a * 3.0
+
+    grad1 = Grad([y1], [a])
+    grad2 = Grad([y2], [a])
+
+    with raises(ValueError):
+        _ = Stack([grad1, grad2])

tests/unit/autojac/_transform/test_select.py

@@ -1,4 +1,8 @@
+from contextlib import nullcontext as does_not_raise
+
 import torch
+from pytest import mark, raises
+from unit._utils import ExceptionContext
 
 from torchjd.autojac._transform import Select, TensorDict
 
@@ -53,3 +57,28 @@ def test_conjunction_of_selects_is_select():
     expected_output = select(input)
 
     assert_tensor_dicts_are_close(output, expected_output)
+
+
+@mark.parametrize(
+    ["key_indices", "required_key_indices", "expectation"],
+    [
+        ([0], [0, 1], does_not_raise()),
+        ([0], [1], raises(ValueError)),
+        ([0, 1], [0], raises(ValueError)),
+        ([], [0], does_not_raise()),
+    ],
+)
+def test_keys_check(
+    key_indices: list[int], required_key_indices: list[int], expectation: ExceptionContext
+):
+    """
+    Tests that the Select transform correctly checks that the keys are a subset of the required
+    keys.
+    """
+
+    all_keys = [torch.tensor(i) for i in range(2)]
+    keys = [all_keys[i] for i in key_indices]
+    required_keys = [all_keys[i] for i in required_key_indices]
+
+    with expectation:
+        _ = Select(keys, required_keys)