Add tests for masking.

tests/update_mask_test.py

@@ -0,0 +1,28 @@
+import typing
+import pytest
+import torch
+from grassmann_tensor import GrassmannTensor
+
+Initialization = tuple[tuple[bool, ...], tuple[tuple[int, int], ...], torch.Tensor]
+
+
+@pytest.fixture(params=[
+    (
+        (True, False),
+        ((2, 2), (3, 1)),
+        torch.randn([4, 4], dtype=torch.float32),
+    ),
+    (
+        (False, True, False),
+        ((1, 1), (2, 2), (1, 1)),
+        torch.randn([2, 4, 2], dtype=torch.float32),
+    ),
+])
+def x(request: typing.Any) -> Initialization:
+    return request.param
+
+
+def test_update_mask(x: Initialization) -> None:
+    tensor = GrassmannTensor(*x)
+    updated_tensor = tensor.update_mask()
+    assert torch.all(updated_tensor.tensor == torch.where(updated_tensor.mask, 0, x[2]))