USTC-KnowledgeComputingLab · hzhangxyz · Jul 30, 2025 · Jul 30, 2025 · Copilot · Jul 30, 2025
diff --git a/parity_tensor/parity_tensor.py b/parity_tensor/parity_tensor.py
@@ -7,6 +7,7 @@
 __all__ = ["ParityTensor"]
 
 import dataclasses
+import functools
 import typing
 import torch
 
@@ -20,11 +21,29 @@ class ParityTensor:
 
     edges: tuple[tuple[int, int], ...]
     tensor: torch.Tensor
+    mask: torch.Tensor | None = None
 
     def __post_init__(self) -> None:
         assert len(self.edges) == self.tensor.dim(), f"Edges length ({len(self.edges)}) must match tensor dimensions ({self.tensor.dim()})."
         for dim, (even, odd) in zip(self.tensor.shape, self.edges):
             assert even >= 0 and odd >= 0 and dim == even + odd, f"Dimension {dim} must equal sum of even ({even}) and odd ({odd}) parts, and both must be non-negative."
+        if self.mask is None:
+            self.mask = self._tensor_mask()
+
+    @classmethod
+    def _unqueeze(cls, tensor: torch.Tensor, index: int, dim: int) -> torch.Tensor:
+        return tensor.view([-1 if i == index else 1 for i in range(dim)])
-        return tensor.view([-1 if i == index else 1 for i in range(dim)])
+        shape = [-1 if i == index else 1 for i in range(dim)]
+        total_elements = tensor.numel()
+        inferred_dim = total_elements // functools.reduce(lambda x, y: x * y, (s for s in shape if s != -1), 1)
+        shape = [inferred_dim if s == -1 else s for s in shape]
+        if total_elements != functools.reduce(lambda x, y: x * y, shape, 1):
+            raise ValueError(f"Cannot reshape tensor of size {total_elements} into shape {shape}.")
+        return tensor.view(shape)
-        return tensor.view([-1 if i == index else 1 for i in range(dim)])
+        shape = [-1 if i == index else 1 for i in range(dim)]
+        total_elements = tensor.numel()
+        inferred_dim = total_elements // functools.reduce(lambda x, y: x * y, (s for s in shape if s != -1), 1)
+        shape = [inferred_dim if s == -1 else s for s in shape]
+        if total_elements != functools.reduce(lambda x, y: x * y, shape, 1):
+            raise ValueError(f"Cannot reshape tensor of size {total_elements} into shape {shape}.")
+        return tensor.view(shape)
+
+    @classmethod
+    def _edge_mask(cls, even: int, odd: int) -> torch.Tensor:
+        return torch.cat([torch.zeros(even, dtype=torch.bool), torch.ones(odd, dtype=torch.bool)])
+
+    def _tensor_mask(self) -> torch.Tensor:
+        return functools.reduce(
+            torch.logical_xor,
+            (self._unqueeze(self._edge_mask(even, odd), index, self.tensor.dim()) for index, (even, odd) in enumerate(self.edges)),
+            torch.ones_like(self.tensor, dtype=torch.bool),
+        )
-        return functools.reduce(
-            torch.logical_xor,
-            (self._unqueeze(self._edge_mask(even, odd), index, self.tensor.dim()) for index, (even, odd) in enumerate(self.edges)),
-            torch.ones_like(self.tensor, dtype=torch.bool),
-        )
+        # Create a grid of indices for each dimension
+        grids = torch.meshgrid(
+            [torch.arange(even + odd) for even, odd in self.edges], indexing="ij"
+        )
+        # Generate masks for each dimension
+        masks = [
+            (grid >= even) for grid, (even, odd) in zip(grids, self.edges)
+        ]
+        # Combine masks using logical XOR
+        return functools.reduce(torch.logical_xor, masks)
-        return functools.reduce(
-            torch.logical_xor,
-            (self._unqueeze(self._edge_mask(even, odd), index, self.tensor.dim()) for index, (even, odd) in enumerate(self.edges)),
-            torch.ones_like(self.tensor, dtype=torch.bool),
-        )
+        # Create a grid of indices for each dimension
+        grids = torch.meshgrid(
+            [torch.arange(even + odd) for even, odd in self.edges], indexing="ij"
+        )
+        # Generate masks for each dimension
+        masks = [
+            (grid >= even) for grid, (even, odd) in zip(grids, self.edges)
+        ]
+        # Combine masks using logical XOR
+        return functools.reduce(torch.logical_xor, masks)
 
     def _validate_edge_compatibility(self, other: ParityTensor) -> None:
         """