diff --git a/grassmann_tensor/tensor.py b/grassmann_tensor/tensor.py index 0236dfc..16160ee 100644 --- a/grassmann_tensor/tensor.py +++ b/grassmann_tensor/tensor.py @@ -205,8 +205,8 @@ def _reorder_indices( torch.zeros([], dtype=torch.bool, device=self.tensor.device), ) flatten_parity = parity.flatten() - even = (~flatten_parity).nonzero().squeeze() - odd = flatten_parity.nonzero().squeeze() + even = (~flatten_parity).nonzero().squeeze(-1) + odd = flatten_parity.nonzero().squeeze(-1) reorder = torch.cat([even, odd], dim=0) total = functools.reduce( @@ -244,8 +244,6 @@ def reshape(self, new_shape: tuple[int | tuple[int, int], ...]) -> GrassmannTens # 5. Apply the sign for merging # 6. Reorder the indices for merging - # pylint: disable=too-many-branches, too-many-locals, too-many-statements - arrow: list[bool] = [] edges: list[tuple[int, int]] = [] shape: list[int] = [] @@ -257,7 +255,7 @@ def reshape(self, new_shape: tuple[int | tuple[int, int], ...]) -> GrassmannTens cursor_plan: int = 0 cursor_self: int = 0 - while True: + while cursor_plan != len(new_shape) or cursor_self != self.tensor.dim(): if new_shape[cursor_plan] == -1: # Does not change arrow.append(self.arrow[cursor_self]) @@ -265,91 +263,132 @@ def reshape(self, new_shape: tuple[int | tuple[int, int], ...]) -> GrassmannTens shape.append(self.tensor.shape[cursor_self]) cursor_self += 1 cursor_plan += 1 - else: - cursor_new_shape = new_shape[cursor_plan] - total = ( - cursor_new_shape - if isinstance(cursor_new_shape, int) - else cursor_new_shape[0] + cursor_new_shape[1] - ) - if total >= self.tensor.shape[cursor_self]: - # Merging - new_cursor_self = cursor_self - self_total = 1 - while True: - self_total *= self.tensor.shape[new_cursor_self] - new_cursor_self += 1 - if self_total == total: - break - assert self_total < total, ( - f"Dimension mismatch with edges {self.edges} and new shape {new_shape}." - ) - assert new_cursor_self < self.tensor.dim(), ( - f"New shape {new_shape} exceeds tensor dimensions {self.tensor.dim()}." - ) - even, odd, reorder, sign = self._reorder_indices( - self.edges[cursor_self:new_cursor_self] - ) - if isinstance(cursor_new_shape, tuple): - assert (even, odd) == cursor_new_shape, ( - f"New even and odd number mismatch during merging {self.edges} to {new_shape}." - ) + continue + if new_shape[cursor_plan] == (1, 0): + # An trivial plan edge + arrow.append(False) + edges.append((1, 0)) + shape.append(1) + cursor_plan += 1 + continue + if self.edges[cursor_self] == (1, 0): + # An trivial self edge + cursor_self += 1 + continue + cursor_new_shape = new_shape[cursor_plan] + total = ( + cursor_new_shape + if isinstance(cursor_new_shape, int) + else cursor_new_shape[0] + cursor_new_shape[1] + ) + # one of total and shape[cursor_self] is not trivial, otherwise it should be handled before + if total == self.tensor.shape[cursor_self]: + # We do not know whether it is merging or splitting, check more + if isinstance(cursor_new_shape, int) or cursor_new_shape == self.edges[cursor_self]: + # If the new shape is exactly the same as the current edge, we treat it as no change arrow.append(self.arrow[cursor_self]) - assert all( - self_arrow == arrow[-1] - for self_arrow in self.arrow[cursor_self:new_cursor_self] - ), ( - f"Cannot merge edges with different arrows {self.arrow[cursor_self:new_cursor_self]}." - ) - edges.append((even, odd)) - shape.append(total) - if cursor_self + 1 != new_cursor_self: - # Really something merged - merging_sign.append((cursor_plan, sign)) - merging_reorder.append((cursor_plan, reorder)) - cursor_self = new_cursor_self + edges.append(self.edges[cursor_self]) + shape.append(self.tensor.shape[cursor_self]) + cursor_self += 1 cursor_plan += 1 - else: - # Splitting - new_cursor_plan = cursor_plan - plan_total = 1 - while True: - new_cursor_new_shape = new_shape[new_cursor_plan] - assert isinstance(new_cursor_new_shape, tuple), ( - f"New shape must be a pair when splitting, got {new_cursor_new_shape}." + continue + # Let's see if there are (0, 1) edges in the remaining self edges, if yes, we treat it as merging, otherwise splitting + cursor_self_finding = cursor_self + cursor_self_found = False + while True: + cursor_self_finding += 1 + if cursor_self_finding == self.tensor.dim(): + break + if self.edges[cursor_self_finding] == (1, 0): + continue + if self.edges[cursor_self_finding] == (0, 1): + cursor_self_found = True + break + break + merging = cursor_self_found + if total > self.tensor.shape[cursor_self]: + merging = True + if total < self.tensor.shape[cursor_self]: + merging = False + if merging: + # Merging between [cursor_self, new_cursor_self) and the another side contains dimension as self_total + new_cursor_self = cursor_self + self_total = 1 + while True: + # Try to include more dimension from self + self_total *= self.tensor.shape[new_cursor_self] + new_cursor_self += 1 + # One dimension included, check if we can stop + if self_total == total: + even, odd, reorder, sign = self._reorder_indices( + self.edges[cursor_self:new_cursor_self] ) - plan_total *= new_cursor_new_shape[0] + new_cursor_new_shape[1] - new_cursor_plan += 1 - if plan_total == self.tensor.shape[cursor_self]: + if isinstance(cursor_new_shape, tuple): + if (even, odd) == cursor_new_shape: + break + else: break - assert plan_total < self.tensor.shape[cursor_self], ( - f"Dimension mismatch with edges {self.edges} and new shape {new_shape}." - ) - assert new_cursor_plan < len(new_shape), ( - f"New shape {new_shape} exceeds specified dimensions {len(new_shape)}." - ) - # new_shape has been verified to be tuple[int, int] in the loop - even, odd, reorder, sign = self._reorder_indices( - typing.cast( - tuple[tuple[int, int], ...], new_shape[cursor_plan:new_cursor_plan] + # For some reason we cannot stop here, continue to include more dimension, check something before continue + assert self_total <= total, ( + f"Dimension mismatch in merging with edges {self.edges} and new shape {new_shape}." + ) + assert new_cursor_self < self.tensor.dim(), ( + f"New shape exceeds in merging with edges {self.edges} and new shape {new_shape}." + ) + # The merging block [cursor_self, new_cursor_self) has been determined + arrow.append(self.arrow[cursor_self]) + assert all( + self_arrow == arrow[-1] + for self_arrow in self.arrow[cursor_self:new_cursor_self] + ), ( + f"Cannot merge edges with different arrows {self.arrow[cursor_self:new_cursor_self]}." + ) + edges.append((even, odd)) + shape.append(total) + merging_sign.append((cursor_plan, sign)) + merging_reorder.append((cursor_plan, reorder)) + cursor_self = new_cursor_self + cursor_plan += 1 + else: + # Splitting between [cursor_plan, new_cursor_plan) and the another side contains dimension as plan_total + new_cursor_plan = cursor_plan + plan_total = 1 + while True: + # Try to include more dimension from new_shape + new_cursor_new_shape = new_shape[new_cursor_plan] + assert isinstance(new_cursor_new_shape, tuple), ( + f"New shape must be a pair when splitting, got {new_cursor_new_shape}." + ) + plan_total *= new_cursor_new_shape[0] + new_cursor_new_shape[1] + new_cursor_plan += 1 + # One dimension included, check if we can stop + if plan_total == self.tensor.shape[cursor_self]: + # new_shape block has been verified to be always tuple[int, int] before + even, odd, reorder, sign = self._reorder_indices( + typing.cast( + tuple[tuple[int, int], ...], new_shape[cursor_plan:new_cursor_plan] + ) ) + if (even, odd) == self.edges[cursor_self]: + break + # For some reason we cannot stop here, continue to include more dimension, check something before continue + assert plan_total <= self.tensor.shape[cursor_self], ( + f"Dimension mismatch in splitting with edges {self.edges} and new shape {new_shape}." ) - assert (even, odd) == self.edges[cursor_self], ( - f"New even and odd number mismatch during splitting {self.edges[cursor_self]} to {new_shape[cursor_plan:new_cursor_plan]}." + assert new_cursor_plan < len(new_shape), ( + f"New shape exceeds in splitting with edges {self.edges} and new shape {new_shape}." ) - for i in range(cursor_plan, new_cursor_plan): - # new_shape has been verified to be tuple[int, int] in the loop - new_cursor_new_shape = typing.cast(tuple[int, int], new_shape[i]) - arrow.append(self.arrow[cursor_self]) - edges.append(new_cursor_new_shape) - shape.append(new_cursor_new_shape[0] + new_cursor_new_shape[1]) - splitting_reorder.append((cursor_self, reorder)) - splitting_sign.append((cursor_self, sign)) - cursor_self += 1 - cursor_plan = new_cursor_plan - - if cursor_plan == len(new_shape) and cursor_self == self.tensor.dim(): - break + # The splitting block [cursor_plan, new_cursor_plan) has been determined + for i in range(cursor_plan, new_cursor_plan): + # new_shape block has been verified to be always tuple[int, int] in the loop + new_cursor_new_shape = typing.cast(tuple[int, int], new_shape[i]) + arrow.append(self.arrow[cursor_self]) + edges.append(new_cursor_new_shape) + shape.append(new_cursor_new_shape[0] + new_cursor_new_shape[1]) + splitting_reorder.append((cursor_self, reorder)) + splitting_sign.append((cursor_self, sign)) + cursor_self += 1 + cursor_plan = new_cursor_plan tensor = self.tensor diff --git a/tests/reshape_test.py b/tests/reshape_test.py index 60f693c..52c49a9 100644 --- a/tests/reshape_test.py +++ b/tests/reshape_test.py @@ -1,3 +1,4 @@ +import random import pytest import torch from grassmann_tensor import GrassmannTensor @@ -27,34 +28,74 @@ def test_reshape_consistency(arrow: tuple[bool, ...], plan_range: tuple[int, int assert torch.allclose(a.tensor, c.tensor) -def test_reshape_merging_dimension_mismatch_edges() -> None: +def insert_trivial_between_elements( + input_list: list[tuple[int, int] | int], p: float +) -> list[tuple[int, int] | int]: + result: list[tuple[int, int] | int] = [] + for i in input_list: + if random.random() < p: + result.append((1, 0)) + result.append(i) + return result + + +@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] # noqa: E741 + 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_trivial_edges(arrow: tuple[bool, ...], plan_range: tuple[int, int]) -> None: + l, h = plan_range # noqa: E741 + 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(insert_trivial_between_elements([-1] * l + [4 ** (h - l)] + [-1] * (5 - h), 0.5)) + b = a.reshape(plan) + c = b.reshape(a.edges) + assert a.edges == c.edges + + +def test_reshape_merging_dimension_mismatch_edges_because_of_nonequal() -> 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"): + with pytest.raises(AssertionError, match="Dimension mismatch in merging"): _ = a.reshape((16, -1, -1)) +def test_reshape_merging_dimension_mismatch_edges_because_of_different_even_odd() -> None: + arrow = (True, True, True, True, True) + edges = ((0, 1), (1, 3), (1, 3), (0, 1), (2, 2)) + a = GrassmannTensor(arrow, edges, torch.randn([1, 4, 4, 1, 4])) + _ = a.reshape((16, -1, -1)) + _ = a.reshape(((6, 10), -1, -1)) + _ = a.reshape(((10, 6), -1)) + _ = a.reshape((4, -1, -1, -1)) + _ = a.reshape(((3, 1), -1, -1, -1)) + with pytest.raises(AssertionError, match="Dimension mismatch in merging"): + _ = a.reshape(((2, 2), -1, -1, -1)) + with pytest.raises(AssertionError, match="Dimension mismatch in merging"): + _ = a.reshape(((1, 3), -1, -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"): + with pytest.raises(AssertionError, match="New shape exceeds in merging"): _ = 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)) @@ -72,27 +113,65 @@ def test_reshape_splitting_shape_type() -> None: _ = a.reshape((2, (2, 2))) -def test_reshape_splitting_dimension_mismatch_edges() -> None: +def test_reshape_splitting_dimension_mismatch_edges_because_of_nonequal() -> 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"): + with pytest.raises(AssertionError, match="Dimension mismatch in splitting"): _ = a.reshape(((4, 4), (2, 2))) +def test_reshape_splitting_dimension_mismatch_edges_because_of_different_even_odd() -> None: + arrow = (True, True) + edges = ((3, 1), (2, 2)) + a = GrassmannTensor(arrow, edges, torch.randn([4, 4])) + _ = a.reshape(((0, 1), (3, 1), (0, 1), (2, 2))) + with pytest.raises(AssertionError, match="Dimension mismatch in splitting"): + _ = a.reshape(((0, 1), (2, 2), (0, 1), (2, 2))) + with pytest.raises(AssertionError, match="Dimension mismatch in splitting"): + _ = a.reshape(((0, 1), (3, 1), (2, 2))) + + def test_reshape_splitting_shape_exceeds() -> None: arrow = (False,) + edges = ((8, 8),) + a = GrassmannTensor(arrow, edges, torch.randn([16])) + with pytest.raises(AssertionError, match="New shape exceeds in splitting"): + _ = a.reshape(((1, 1), (1, 1))) + + +def test_reshape_equal_edges_trivial() -> None: + arrow = (True,) edges = ((2, 2),) a = GrassmannTensor(arrow, edges, torch.randn([4])) - with pytest.raises(AssertionError, match="exceeds specified dimensions"): - _ = a.reshape(((3, 0), (0, 1))) + _ = a.reshape((4,)) + _ = a.reshape(((2, 2),)) -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))) +def test_reshape_equal_edges_nontrivial_splitting() -> None: + arrow = (True,) + edges = ((1, 3),) + a = GrassmannTensor(arrow, edges, torch.randn([4])) + _ = a.reshape(((3, 1), (1, 0), (0, 1))) + + +def test_reshape_equal_edges_nontrivial_splitting_with_other_edge() -> None: + arrow = (True, True) + edges = ((1, 3), (2, 2)) + a = GrassmannTensor(arrow, edges, torch.randn([4, 4])) + _ = a.reshape(((3, 1), (1, 0), (0, 1), (2, 2))) + + +def test_reshape_equal_edges_nontrivial_merging() -> None: + arrow = (True, True, True) + edges = ((1, 3), (1, 0), (0, 1)) + a = GrassmannTensor(arrow, edges, torch.randn([4, 1, 1])) + _ = a.reshape(((3, 1),)) + + +def test_reshape_equal_edges_nontrivial_merging_with_other_edge() -> None: + arrow = (True, True, True, True) + edges = ((1, 3), (1, 0), (0, 1), (2, 2)) + a = GrassmannTensor(arrow, edges, torch.randn([4, 1, 1, 4])) + _ = a.reshape(((3, 1), (2, 2)))