diff --git a/tests/reshape_test.py b/tests/reshape_test.py index 6d32e90..f04b273 100644 --- a/tests/reshape_test.py +++ b/tests/reshape_test.py @@ -15,3 +15,74 @@ def test_reshape_consistency(arrow: tuple[bool, ...], plan_range: tuple[int, int b = a.reshape(plan) c = b.reshape(a.edges) assert torch.allclose(a.tensor, c.tensor) + + +def test_reshape_merging_dimension_mismatch_edges() -> None: + arrow = (True, True, True) + edges = ((2, 2), (8, 8), (2, 2)) + a = GrassmannTensor(arrow, edges, torch.randn([4, 16, 4])) + _ = a.reshape((64, -1)) + _ = a.reshape((-1, 64)) + with pytest.raises(AssertionError, match="Dimension mismatch with edges"): + _ = a.reshape((16, -1, -1)) + + +def test_reshape_merging_new_shape_exceeds() -> None: + arrow = (True,) + edges = ((2, 2),) + a = GrassmannTensor(arrow, edges, torch.randn([4])) + with pytest.raises(AssertionError, match="exceeds tensor dimensions"): + _ = a.reshape((16, -1)) + + +def test_reshape_merging_even_odd_mismatch() -> None: + arrow = (True, True, True) + edges = ((2, 2), (8, 8), (2, 2)) + a = GrassmannTensor(arrow, edges, torch.randn([4, 16, 4])) + _ = a.reshape(((32, 32), (2, 2))) + _ = a.reshape(((2, 2), (32, 32))) + with pytest.raises(AssertionError, match="New even and odd number mismatch during merging"): + _ = a.reshape(((30, 34), (2, 2))) + + +def test_reshape_merging_mixed_arrows() -> None: + arrow = (True, False, True) + edges = ((2, 2), (2, 2), (2, 2)) + a = GrassmannTensor(arrow, edges, torch.randn([4, 4, 4])) + with pytest.raises(AssertionError, match="Cannot merge edges with different arrows"): + _ = a.reshape((64,)) + + +def test_reshape_splitting_shape_type() -> None: + arrow = (True,) + edges = ((8, 8),) + a = GrassmannTensor(arrow, edges, torch.randn([16])) + _ = a.reshape(((2, 2), (2, 2))) + with pytest.raises(AssertionError, match="New shape must be a pair when splitting"): + _ = a.reshape((2, (2, 2))) + + +def test_reshape_splitting_dimension_mismatch_edges() -> None: + arrow = (True,) + edges = ((8, 8),) + a = GrassmannTensor(arrow, edges, torch.randn([16])) + _ = a.reshape(((2, 2), (2, 2))) + with pytest.raises(AssertionError, match="Dimension mismatch with edges"): + _ = a.reshape(((4, 4), (2, 2))) + + +def test_reshape_splitting_shape_exceeds() -> None: + arrow = (False,) + edges = ((2, 2),) + a = GrassmannTensor(arrow, edges, torch.randn([4])) + with pytest.raises(AssertionError, match="exceeds specified dimensions"): + _ = a.reshape(((3, 0), (0, 1))) + + +def test_reshape_splitting_even_odd_mismatch() -> None: + arrow = (False,) + edges = ((6, 10),) + a = GrassmannTensor(arrow, edges, torch.randn([16])) + _ = a.reshape(((1, 3), (3, 1))) + with pytest.raises(AssertionError, match="New even and odd number mismatch during splitting"): + _ = a.reshape(((2, 2), (2, 2)))