fix(reshape): support reshaping 1D Grassmann tensors to all-ones shapes

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))

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)))