From f0ce9752524cfd38cda5309c1cac756de24ceb1f Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Sun, 10 Aug 2025 22:45:02 +0800 Subject: [PATCH] Add a simple test for reshape function. --- tests/reshape_test.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 tests/reshape_test.py diff --git a/tests/reshape_test.py b/tests/reshape_test.py new file mode 100644 index 0000000..6d32e90 --- /dev/null +++ b/tests/reshape_test.py @@ -0,0 +1,17 @@ +import pytest +import torch +from grassmann_tensor import GrassmannTensor + + +@pytest.mark.parametrize("arrow", [(i, j, k, l, m) for i in [False, True] for j in [False, True] for k in [False, True] for l in [False, True] for m in [False, True]]) +@pytest.mark.parametrize("plan_range", [(i, j) for i in range(5) for j in range(5) if j > i]) +def test_reshape_consistency(arrow: tuple[bool, ...], plan_range: tuple[int, int]) -> None: + l, h = plan_range + if not all(arrow[l:h]) and any(arrow[l:h]): + pytest.skip("Invalid reshape plan for the given arrow configuration.") + edge = (2, 2) + a = GrassmannTensor(arrow, (edge, edge, edge, edge, edge), torch.randn([4, 4, 4, 4, 4])) + plan = tuple([-1] * l + [4**(h - l)] + [-1] * (5 - h)) + b = a.reshape(plan) + c = b.reshape(a.edges) + assert torch.allclose(a.tensor, c.tensor)