test(autogram): Improve code quality (#464)

* Extract rng forking into contexts.py
* Make _forward_pass do rng forking
* Make _forward_pass take reduction parameter
* Make forward_pass public
* Use forward_pass in test_engine.py, stop reseeding (it's now done by forward_pass)
* Make zipping strict in make_mse_loss_fn
* Stop requiring params in autograd_gramian_forward_backward
* Improve parameter order of autogram_forward_backward
* Rename some variables
* Factorize input and target creation into make_inputs_and_targets
* Reorder some code

Author: ValerianRey

tests/speed/autogram/grad_vs_jac_vs_gram.py

@@ -14,7 +14,6 @@
     NoFreeParam,
     SqueezeNet,
     WithTransformerLarge,
-    get_in_out_shapes,
 )
 from utils.forward_backwards import (
     autograd_forward_backward,
@@ -23,7 +22,7 @@
     autojac_forward_backward,
     make_mse_loss_fn,
 )
-from utils.tensors import make_tensors
+from utils.tensors import make_inputs_and_targets
 
 from torchjd.aggregation import Mean
 from torchjd.autogram import Engine
@@ -43,9 +42,7 @@
 
 def compare_autograd_autojac_and_autogram_speed(factory: ModuleFactory, batch_size: int):
     model = factory()
-    input_shapes, output_shapes = get_in_out_shapes(model)
-    inputs = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model, batch_size)
     loss_fn = make_mse_loss_fn(targets)
 
     A = Mean()
@@ -64,7 +61,7 @@ def init_fn_autograd():
         fn_autograd()
 
     def fn_autograd_gramian():
-        autograd_gramian_forward_backward(model, inputs, list(model.parameters()), loss_fn, W)
+        autograd_gramian_forward_backward(model, inputs, loss_fn, W)
 
     def init_fn_autograd_gramian():
         torch.cuda.empty_cache()
@@ -80,7 +77,7 @@ def init_fn_autojac():
         fn_autojac()
 
     def fn_autogram():
-        autogram_forward_backward(model, engine, W, inputs, loss_fn)
+        autogram_forward_backward(model, inputs, loss_fn, engine, W)
 
     def init_fn_autogram():
         torch.cuda.empty_cache()

tests/unit/autogram/test_engine.py

@@ -61,21 +61,21 @@
     WithSomeFrozenModule,
     WithTransformer,
     WithTransformerLarge,
-    get_in_out_shapes,
 )
 from utils.dict_assertions import assert_tensor_dicts_are_close
 from utils.forward_backwards import (
     autograd_forward_backward,
     autogram_forward_backward,
     compute_gramian,
     compute_gramian_with_autograd,
+    forward_pass,
     make_mse_loss_fn,
     reduce_to_first_tensor,
     reduce_to_matrix,
     reduce_to_scalar,
     reduce_to_vector,
 )
-from utils.tensors import make_tensors, ones_, randn_, zeros_
+from utils.tensors import make_inputs_and_targets, ones_, randn_, zeros_
 
 from torchjd.aggregation import UPGradWeighting
 from torchjd.autogram._engine import Engine
@@ -144,22 +144,14 @@ def _assert_gramian_is_equivalent_to_autograd(
     factory: ModuleFactory, batch_size: int, batch_dim: int | None
 ):
     model_autograd, model_autogram = factory(), factory()
-    input_shapes, output_shapes = get_in_out_shapes(model_autograd)
-
     engine = Engine(model_autogram, batch_dim=batch_dim)
-
-    inputs = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model_autograd, batch_size)
     loss_fn = make_mse_loss_fn(targets)
 
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_autograd(inputs)
-    losses = reduce_to_vector(loss_fn(output))
+    losses = forward_pass(model_autograd, inputs, loss_fn, reduce_to_vector)
     autograd_gramian = compute_gramian_with_autograd(losses, list(model_autograd.parameters()))
 
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_autogram(inputs)
-    losses = reduce_to_vector(loss_fn(output))
+    losses = forward_pass(model_autogram, inputs, loss_fn, reduce_to_vector)
     autogram_gramian = engine.compute_gramian(losses)
 
     assert_close(autogram_gramian, autograd_gramian, rtol=1e-4, atol=3e-5)
@@ -255,26 +247,18 @@ def test_compute_gramian_various_output_shapes(
 
     factory = ModuleFactory(Ndim2Output)
     model_autograd, model_autogram = factory(), factory()
-    input_shapes, output_shapes = get_in_out_shapes(model_autograd)
-
-    engine = Engine(model_autogram, batch_dim=batch_dim)
-
-    inputs = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model_autograd, batch_size)
     loss_fn = make_mse_loss_fn(targets)
 
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_autograd(inputs)
-    losses = reduction(loss_fn(output))
+    losses = forward_pass(model_autograd, inputs, loss_fn, reduction)
     reshaped_losses = torch.movedim(losses, movedim_source, movedim_destination)
     # Go back to a vector so that compute_gramian_with_autograd works
     loss_vector = reshaped_losses.reshape([-1])
     autograd_gramian = compute_gramian_with_autograd(loss_vector, list(model_autograd.parameters()))
     expected_gramian = reshape_gramian(autograd_gramian, list(reshaped_losses.shape))
 
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_autogram(inputs)
-    losses = reduction(loss_fn(output))
+    engine = Engine(model_autogram, batch_dim=batch_dim)
+    losses = forward_pass(model_autogram, inputs, loss_fn, reduction)
     reshaped_losses = torch.movedim(losses, movedim_source, movedim_destination)
     autogram_gramian = engine.compute_gramian(reshaped_losses)
 
@@ -296,30 +280,20 @@ def test_compute_partial_gramian(gramian_module_names: set[str], batch_dim: int
     the model parameters is specified.
     """
 
-    factory = ModuleFactory(SimpleBranched)
-    model = factory()
-    input_shapes, output_shapes = get_in_out_shapes(model)
+    model = SimpleBranched()
     batch_size = 64
-
-    input = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model, batch_size)
     loss_fn = make_mse_loss_fn(targets)
-
-    output = model(input)
-    losses = reduce_to_vector(loss_fn(output))
-
     gramian_modules = [model.get_submodule(name) for name in gramian_module_names]
     gramian_params = []
     for m in gramian_modules:
         gramian_params += list(m.parameters())
 
+    losses = forward_pass(model, inputs, loss_fn, reduce_to_vector)
     autograd_gramian = compute_gramian_with_autograd(losses, gramian_params, retain_graph=True)
-    torch.manual_seed(0)
 
     engine = Engine(*gramian_modules, batch_dim=batch_dim)
-
-    output = model(input)
-    losses = reduce_to_vector(loss_fn(output))
+    losses = forward_pass(model, inputs, loss_fn, reduce_to_vector)
     gramian = engine.compute_gramian(losses)
 
     assert_close(gramian, autograd_gramian)
@@ -331,22 +305,15 @@ def test_iwrm_steps_with_autogram(factory: ModuleFactory, batch_size: int, batch
     """Tests that the autogram engine doesn't raise any error during several IWRM iterations."""
 
     n_iter = 3
-
     model = factory()
-    input_shapes, output_shapes = get_in_out_shapes(model)
-
     weighting = UPGradWeighting()
-
     engine = Engine(model, batch_dim=batch_dim)
     optimizer = SGD(model.parameters(), lr=1e-7)
 
     for i in range(n_iter):
-        inputs = make_tensors(batch_size, input_shapes)
-        targets = make_tensors(batch_size, output_shapes)
+        inputs, targets = make_inputs_and_targets(model, batch_size)
         loss_fn = make_mse_loss_fn(targets)
-
-        autogram_forward_backward(model, engine, weighting, inputs, loss_fn)
-
+        autogram_forward_backward(model, inputs, loss_fn, engine, weighting)
         optimizer.step()
         model.zero_grad()
 
@@ -363,29 +330,22 @@ def test_autograd_while_modules_are_hooked(
     """
 
     model, model_autogram = factory(), factory()
-    input_shapes, output_shapes = get_in_out_shapes(model)
-
-    input = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model, batch_size)
     loss_fn = make_mse_loss_fn(targets)
 
-    torch.manual_seed(0)  # Fix randomness for random models
-    autograd_forward_backward(model, input, loss_fn)
+    autograd_forward_backward(model, inputs, loss_fn)
     autograd_grads = {name: p.grad for name, p in model.named_parameters() if p.grad is not None}
 
     # Hook modules and optionally compute the Gramian
     engine = Engine(model_autogram, batch_dim=batch_dim)
     if use_engine:
-        torch.manual_seed(0)  # Fix randomness for random models
-        output = model_autogram(input)
-        losses = reduce_to_vector(loss_fn(output))
+        losses = forward_pass(model_autogram, inputs, loss_fn, reduce_to_vector)
         _ = engine.compute_gramian(losses)
 
     # Verify that even with the hooked modules, autograd works normally when not using the engine.
     # Results should be the same as a normal call to autograd, and no time should be spent computing
     # the gramian at all.
-    torch.manual_seed(0)  # Fix randomness for random models
-    autograd_forward_backward(model_autogram, input, loss_fn)
+    autograd_forward_backward(model_autogram, inputs, loss_fn)
     grads = {name: p.grad for name, p in model_autogram.named_parameters() if p.grad is not None}
 
     assert_tensor_dicts_are_close(grads, autograd_grads)
@@ -416,12 +376,11 @@ def test_compute_gramian_manual():
 
     in_dims = 18
     out_dims = 25
-
     factory = ModuleFactory(Linear, in_dims, out_dims)
     model = factory()
-    engine = Engine(model, batch_dim=None)
-
     input = randn_(in_dims)
+
+    engine = Engine(model, batch_dim=None)
     output = model(input)
     gramian = engine.compute_gramian(output)
 
@@ -462,21 +421,19 @@ def test_reshape_equivariance(shape: list[int]):
 
     input_size = shape[0]
     output_size = prod(shape[1:])
-
     factory = ModuleFactory(Linear, input_size, output_size)
     model1, model2 = factory(), factory()
+    input = randn_([input_size])
 
     engine1 = Engine(model1, batch_dim=None)
-    engine2 = Engine(model2, batch_dim=None)
-
-    input = randn_([input_size])
     output = model1(input)
-    reshaped_output = model2(input).reshape(shape[1:])
-
     gramian = engine1.compute_gramian(output)
-    reshaped_gramian = engine2.compute_gramian(reshaped_output)
     expected_reshaped_gramian = reshape_gramian(gramian, shape[1:])
 
+    engine2 = Engine(model2, batch_dim=None)
+    reshaped_output = model2(input).reshape(shape[1:])
+    reshaped_gramian = engine2.compute_gramian(reshaped_output)
+
     assert_close(reshaped_gramian, expected_reshaped_gramian)
 
 
@@ -502,21 +459,19 @@ def test_movedim_equivariance(shape: list[int], source: list[int], destination:
 
     input_size = shape[0]
     output_size = prod(shape[1:])
-
     factory = ModuleFactory(Linear, input_size, output_size)
     model1, model2 = factory(), factory()
+    input = randn_([input_size])
 
     engine1 = Engine(model1, batch_dim=None)
-    engine2 = Engine(model2, batch_dim=None)
-
-    input = randn_([input_size])
     output = model1(input).reshape(shape[1:])
-    moved_output = model2(input).reshape(shape[1:]).movedim(source, destination)
-
     gramian = engine1.compute_gramian(output)
-    moved_gramian = engine2.compute_gramian(moved_output)
     expected_moved_gramian = movedim_gramian(gramian, source, destination)
 
+    engine2 = Engine(model2, batch_dim=None)
+    moved_output = model2(input).reshape(shape[1:]).movedim(source, destination)
+    moved_gramian = engine2.compute_gramian(moved_output)
+
     assert_close(moved_gramian, expected_moved_gramian)
 
 
@@ -545,18 +500,16 @@ def test_batched_non_batched_equivalence(shape: list[int], batch_dim: int):
     input_size = prod(non_batched_shape)
     batch_size = shape[batch_dim]
     output_size = input_size
-
     factory = ModuleFactory(Linear, input_size, output_size)
     model1, model2 = factory(), factory()
+    input = randn_([batch_size, input_size])
 
     engine1 = Engine(model1, batch_dim=batch_dim)
-    engine2 = Engine(model2, batch_dim=None)
-
-    input = randn_([batch_size, input_size])
     output1 = model1(input).reshape([batch_size] + non_batched_shape).movedim(0, batch_dim)
-    output2 = model2(input).reshape([batch_size] + non_batched_shape).movedim(0, batch_dim)
-
     gramian1 = engine1.compute_gramian(output1)
+
+    engine2 = Engine(model2, batch_dim=None)
+    output2 = model2(input).reshape([batch_size] + non_batched_shape).movedim(0, batch_dim)
     gramian2 = engine2.compute_gramian(output2)
 
     assert_close(gramian1, gramian2)
@@ -573,24 +526,15 @@ def test_batched_non_batched_equivalence_2(factory: ModuleFactory, batch_size: i
     """
 
     model_0, model_none = factory(), factory()
-    input_shapes, output_shapes = get_in_out_shapes(model_0)
-
-    engine_0 = Engine(model_0, batch_dim=0)
-    engine_none = Engine(model_none, batch_dim=None)
-
-    inputs = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model_0, batch_size)
     loss_fn = make_mse_loss_fn(targets)
 
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_0(inputs)
-    losses_0 = reduce_to_vector(loss_fn(output))
-
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_none(inputs)
-    losses_none = reduce_to_vector(loss_fn(output))
-
+    engine_0 = Engine(model_0, batch_dim=0)
+    losses_0 = forward_pass(model_0, inputs, loss_fn, reduce_to_vector)
     gramian_0 = engine_0.compute_gramian(losses_0)
+
+    engine_none = Engine(model_none, batch_dim=None)
+    losses_none = forward_pass(model_none, inputs, loss_fn, reduce_to_vector)
     gramian_none = engine_none.compute_gramian(losses_none)
 
     assert_close(gramian_0, gramian_none, rtol=1e-4, atol=1e-5)

tests/utils/architectures.py

@@ -6,6 +6,7 @@
 from torch import Tensor, nn
 from torch.nn import Flatten, ReLU
 from torch.utils._pytree import PyTree
+from utils.contexts import fork_rng
 
 
 class ModuleFactory:
@@ -15,9 +16,7 @@ def __init__(self, architecture: type[nn.Module], *args, **kwargs):
         self.kwargs = kwargs
 
     def __call__(self) -> nn.Module:
-        devices = [DEVICE] if DEVICE.type == "cuda" else []
-        with torch.random.fork_rng(devices=devices, device_type=DEVICE.type):
-            torch.random.manual_seed(0)
+        with fork_rng(seed=0):
             return self.architecture(*self.args, **self.kwargs).to(device=DEVICE)
 
     def __str__(self) -> str:

tests/utils/contexts.py

@@ -1,4 +1,16 @@
-from contextlib import AbstractContextManager
-from typing import TypeAlias
+from collections.abc import Generator
+from contextlib import AbstractContextManager, contextmanager
+from typing import Any, TypeAlias
+
+import torch
+from device import DEVICE
 
 ExceptionContext: TypeAlias = AbstractContextManager[Exception | None]
+
+
+@contextmanager
+def fork_rng(seed: int = 0) -> Generator[Any, None, None]:
+    devices = [DEVICE] if DEVICE.type == "cuda" else []
+    with torch.random.fork_rng(devices=devices, device_type=DEVICE.type) as ctx:
+        torch.manual_seed(seed)
+        yield ctx