Refactor accumulate tests to use loops and assert helpers

Author: ValerianRey

tests/unit/autojac/_transform/test_accumulate.py

@@ -1,4 +1,5 @@
 from pytest import mark, raises
+from unit.autojac._asserts import assert_grad_close, assert_jac_close
 from utils.dict_assertions import assert_tensor_dicts_are_close
 from utils.tensors import ones_, tensor_, zeros_
 
@@ -11,25 +12,18 @@ def test_single_grad_accumulation():
     once.
     """
 
-    key1 = zeros_([], requires_grad=True)
-    key2 = zeros_([1], requires_grad=True)
-    key3 = zeros_([2, 3], requires_grad=True)
-    value1 = ones_([])
-    value2 = ones_([1])
-    value3 = ones_([2, 3])
-    input = {key1: value1, key2: value2, key3: value3}
+    shapes = [[], [1], [2, 3]]
+    keys = [zeros_(shape, requires_grad=True) for shape in shapes]
+    values = [ones_(shape) for shape in shapes]
+    input = dict(zip(keys, values))
 
     accumulate = AccumulateGrad()
 
     output = accumulate(input)
-    expected_output = {}
+    assert_tensor_dicts_are_close(output, {})
 
-    assert_tensor_dicts_are_close(output, expected_output)
-
-    grads = {key1: key1.grad, key2: key2.grad, key3: key3.grad}
-    expected_grads = {key1: value1, key2: value2, key3: value3}
-
-    assert_tensor_dicts_are_close(grads, expected_grads)
+    for key, value in zip(keys, values):
+        assert_grad_close(key, value)
 
 
 @mark.parametrize("iterations", [1, 2, 4, 10, 13])
@@ -39,28 +33,18 @@ def test_multiple_grad_accumulations(iterations: int):
     `iterations` times.
     """
 
-    key1 = zeros_([], requires_grad=True)
-    key2 = zeros_([1], requires_grad=True)
-    key3 = zeros_([2, 3], requires_grad=True)
-    value1 = ones_([])
-    value2 = ones_([1])
-    value3 = ones_([2, 3])
-
+    shapes = [[], [1], [2, 3]]
+    keys = [zeros_(shape, requires_grad=True) for shape in shapes]
+    values = [ones_(shape) for shape in shapes]
     accumulate = AccumulateGrad()
 
     for i in range(iterations):
         # Clone values to ensure that we accumulate values that are not ever used afterwards
-        input = {key1: value1.clone(), key2: value2.clone(), key3: value3.clone()}
+        input = {key: value.clone() for key, value in zip(keys, values)}
         accumulate(input)
 
-    grads = {key1: key1.grad, key2: key2.grad, key3: key3.grad}
-    expected_grads = {
-        key1: iterations * value1,
-        key2: iterations * value2,
-        key3: iterations * value3,
-    }
-
-    assert_tensor_dicts_are_close(grads, expected_grads)
+    for key, value in zip(keys, values):
+        assert_grad_close(key, iterations * value)
 
 
 def test_accumulate_grad_fails_when_no_requires_grad():
@@ -111,25 +95,18 @@ def test_single_jac_accumulation():
     once.
     """
 
-    key1 = zeros_([], requires_grad=True)
-    key2 = zeros_([1], requires_grad=True)
-    key3 = zeros_([2, 3], requires_grad=True)
-    value1 = ones_([4])
-    value2 = ones_([4, 1])
-    value3 = ones_([4, 2, 3])
-    input = {key1: value1, key2: value2, key3: value3}
+    shapes = [[], [1], [2, 3]]
+    keys = [zeros_(shape, requires_grad=True) for shape in shapes]
+    values = [ones_([4] + shape) for shape in shapes]
+    input = dict(zip(keys, values))
 
     accumulate = AccumulateJac()
 
     output = accumulate(input)
-    expected_output = {}
+    assert_tensor_dicts_are_close(output, {})
 
-    assert_tensor_dicts_are_close(output, expected_output)
-
-    jacs = {key1: key1.jac, key2: key2.jac, key3: key3.jac}
-    expected_jacs = {key1: value1, key2: value2, key3: value3}
-
-    assert_tensor_dicts_are_close(jacs, expected_jacs)
+    for key, value in zip(keys, values):
+        assert_jac_close(key, value)
 
 
 @mark.parametrize("iterations", [1, 2, 4, 10, 13])
@@ -139,28 +116,19 @@ def test_multiple_jac_accumulations(iterations: int):
     `iterations` times.
     """
 
-    key1 = zeros_([], requires_grad=True)
-    key2 = zeros_([1], requires_grad=True)
-    key3 = zeros_([2, 3], requires_grad=True)
-    value1 = ones_([4])
-    value2 = ones_([4, 1])
-    value3 = ones_([4, 2, 3])
+    shapes = [[], [1], [2, 3]]
+    keys = [zeros_(shape, requires_grad=True) for shape in shapes]
+    values = [ones_([4] + shape) for shape in shapes]
 
     accumulate = AccumulateJac()
 
     for i in range(iterations):
         # Clone values to ensure that we accumulate values that are not ever used afterwards
-        input = {key1: value1.clone(), key2: value2.clone(), key3: value3.clone()}
+        input = {key: value.clone() for key, value in zip(keys, values)}
         accumulate(input)
 
-    jacs = {key1: key1.jac, key2: key2.jac, key3: key3.jac}
-    expected_jacs = {
-        key1: iterations * value1,
-        key2: iterations * value2,
-        key3: iterations * value3,
-    }
-
-    assert_tensor_dicts_are_close(jacs, expected_jacs)
+    for key, value in zip(keys, values):
+        assert_jac_close(key, iterations * value)
 
 
 def test_accumulate_jac_fails_when_no_requires_grad():