diff --git a/grassmann_tensor/tensor.py b/grassmann_tensor/tensor.py index fa22531..ae9f49a 100644 --- a/grassmann_tensor/tensor.py +++ b/grassmann_tensor/tensor.py @@ -295,9 +295,28 @@ def reshape(self, new_shape: tuple[int | tuple[int, int], ...]) -> GrassmannTens tensor = self.tensor.reshape(()) return GrassmannTensor(_arrow=(), _edges=(), _tensor=tensor) + if new_shape == (1,) and int(self.tensor.numel()) == 1: + eo = self._calculate_even_odd() + new_shape = (eo,) + cursor_plan: int = 0 cursor_self: int = 0 while cursor_plan != len(new_shape) or cursor_self != self.tensor.dim(): + if cursor_self == self.tensor.dim() and cursor_plan != len(new_shape): + new_shape_check = new_shape[cursor_plan] + if (isinstance(new_shape_check, int) and new_shape_check == 1) or ( + new_shape_check == (1, 0) + ): + arrow.append(False) + edges.append((1, 0)) + shape.append(1) + cursor_plan += 1 + continue + raise AssertionError( + "New shape exceeds after exhausting self dimensions: " + f"edges={self.edges}, new_shape={new_shape}" + ) + if cursor_plan != len(new_shape) and new_shape[cursor_plan] == -1: # Does not change arrow.append(self.arrow[cursor_self]) @@ -306,7 +325,11 @@ def reshape(self, new_shape: tuple[int | tuple[int, int], ...]) -> GrassmannTens cursor_self += 1 cursor_plan += 1 continue - elif cursor_plan != len(new_shape) and new_shape[cursor_plan] == (1, 0): + elif ( + cursor_plan != len(new_shape) + and new_shape[cursor_plan] == (1, 0) + and cursor_plan < len(new_shape) - 1 + ): # A trivial plan edge arrow.append(False) edges.append((1, 0)) diff --git a/tests/reshape_test.py b/tests/reshape_test.py index a1dac6f..2d2d452 100644 --- a/tests/reshape_test.py +++ b/tests/reshape_test.py @@ -197,3 +197,37 @@ def test_reshape_with_none_edge_assertion() -> None: _ = GrassmannTensor((), (), torch.tensor(2333)).reshape((1, -1)) with pytest.raises(AssertionError, match="Ambiguous integer dim"): _ = GrassmannTensor((), (), torch.tensor(2333)).reshape((2, 2)) + + +@pytest.mark.parametrize( + "arrow, edges, tensor", + [ + ((True, True), ((0, 1), (0, 1)), torch.tensor([[2333]])), + ((True, True, True), ((0, 1), (1, 0), (0, 1)), torch.tensor([[[2333]]])), + ], +) +@pytest.mark.parametrize( + "shape", + [ + (1,), + (1, 1), + (1, 1, 1), + (1, 1, 1, 1), + ], +) +def test_reshape_with_one_dimension( + arrow: tuple[bool, ...], + edges: tuple[tuple[int, int], ...], + tensor: torch.Tensor, + shape: tuple[int, ...], +) -> None: + a = GrassmannTensor(arrow, edges, tensor).reshape(shape) + assert ( + len(a.arrow) == len(shape) and len(a.edges) == len(shape) and a.tensor.dim() == len(shape) + ) + + +def test_reshape_trailing_nontrivial_dim_raises() -> None: + a = GrassmannTensor((True,), ((2, 2),), torch.randn([4])) + with pytest.raises(AssertionError, match="New shape exceeds after exhausting self dimensions"): + _ = a.reshape((-1, (2, 2)))