Add test_identity

Co-authored-by: Radka Jersakova <r.jersakova@gmail.com>
Co-authored-by: Marjan Famili <marjanfamili@users.noreply.github.com>
Co-authored-by: Christopher Iliffe Sprague <cisprague@users.noreply.github.com>
Co-authored-by: Paolo Conti <ContiPaolo@users.noreply.github.com>

Author: 5 people

tests/conftest.py

@@ -9,6 +9,28 @@
 from auto_cast.types import Batch, EncodedBatch
 
 
+def _make_batch(
+    batch_size: int = 1,
+    t: int = 1,
+    w: int = 2,
+    h: int = 3,
+    c: int = 2,
+    const_c: int = 1,
+    scalar_c: int = 1,
+) -> Batch:
+    input_fields = torch.arange(batch_size * t * w * h * c, dtype=torch.float32)
+    input_fields = input_fields.view(batch_size, t, w, h, c)
+    output_fields = torch.zeros(batch_size, t, w, h, c)
+    constant_fields = torch.ones(batch_size, w, h, const_c)
+    constant_scalars = torch.full((batch_size, scalar_c), 5.0)
+    return Batch(
+        input_fields=input_fields,
+        output_fields=output_fields,
+        constant_scalars=constant_scalars,
+        constant_fields=constant_fields,
+    )
+
+
 def assert_output_valid(output: Tensor, expected_shape: tuple, name: str = "Output"):
     """Assert output has expected shape and contains no NaN values."""
     assert output.shape == expected_shape, (

tests/encoders/test_identity.py

@@ -0,0 +1,12 @@
+import torch
+from conftest import _make_batch
+
+from auto_cast.encoders.identity import IdentityEncoder
+
+
+def test_identity():
+    batch = _make_batch()
+    encoder = IdentityEncoder()
+    encoded_batch = encoder.encode_batch(batch)
+    assert torch.allclose(encoded_batch.encoded_inputs, batch.input_fields)
+    assert torch.allclose(encoded_batch.encoded_output_fields, batch.output_fields)

tests/encoders/test_permute_concat.py

@@ -1,30 +1,8 @@
 import torch
+from conftest import _make_batch
 from einops import rearrange
 
 from auto_cast.encoders.permute_concat import PermuteConcat
-from auto_cast.types import Batch
-
-
-def _make_batch(
-    batch_size: int = 1,
-    t: int = 1,
-    w: int = 2,
-    h: int = 3,
-    c: int = 2,
-    const_c: int = 1,
-    scalar_c: int = 1,
-) -> Batch:
-    input_fields = torch.arange(batch_size * t * w * h * c, dtype=torch.float32)
-    input_fields = input_fields.view(batch_size, t, w, h, c)
-    output_fields = torch.zeros(batch_size, t, w, h, c)
-    constant_fields = torch.ones(batch_size, w, h, const_c)
-    constant_scalars = torch.full((batch_size, scalar_c), 5.0)
-    return Batch(
-        input_fields=input_fields,
-        output_fields=output_fields,
-        constant_scalars=constant_scalars,
-        constant_fields=constant_fields,
-    )
 
 
 def test_permute_concat_with_constants():