Add the basic parity tensor class and its arithmetic operators.

Author: hzhangxyz

parity_tensor/__init__.py

@@ -2,6 +2,7 @@
 A parity tensor package.
 """
 
-__all__ = ["__version__"]
+__all__ = ["__version__", "ParityTensor"]
 
 from .version import __version__
+from .parity_tensor import ParityTensor

parity_tensor/parity_tensor.py

@@ -0,0 +1,152 @@
+"""
+A parity tensor class.
+"""
+
+from __future__ import annotations
+
+__all__ = ["ParityTensor"]
+
+import dataclasses
+import torch
+
+
+@dataclasses.dataclass
+class ParityTensor:
+    """
+    A parity tensor class, which stores a tensor along with information about its edges.
+    Each dimension of the tensor is composed of an even and an odd part, represented as a pair of integers.
+    """
+
+    edges: tuple[tuple[int, int], ...]
+    tensor: torch.Tensor
+
+    def __post_init__(self):
+        assert len(self.edges) == self.tensor.dim(), "Edges length must match tensor dimensions."
+        for dim, (even, odd) in zip(self.tensor.shape, self.edges):
+            assert even >= 0 and odd >= 0 and dim == even + odd, "Each dimension must match the sum of even and odd parts."
+
+    def _validate_edge_compatibility(self, other: ParityTensor) -> None:
+        """
+        Validate that the edges of two ParityTensor instances are compatible for arithmetic operations.
+        """
+        assert self.edges == other.edges, "Edges must match for arithmetic operations."
+
+    def __add__(self, other):
+        if isinstance(other, ParityTensor):  # pylint: disable=no-else-return
+            self._validate_edge_compatibility(other)
+            return ParityTensor(
+                edges=self.edges,
+                tensor=self.tensor + other.tensor,
+            )
+        else:
+            return ParityTensor(
+                edges=self.edges,
+                tensor=self.tensor + other,
+            )
+
+    def __radd__(self, other):
+        return ParityTensor(
+            edges=self.edges,
+            tensor=other + self.tensor,
+        )
+
+    def __iadd__(self, other):
+        if isinstance(other, ParityTensor):
+            self._validate_edge_compatibility(other)
+            self.tensor += other.tensor
+        else:
+            self.tensor += other
+        return self
+
+    def __sub__(self, other):
+        if isinstance(other, ParityTensor):  # pylint: disable=no-else-return
+            self._validate_edge_compatibility(other)
+            return ParityTensor(
+                edges=self.edges,
+                tensor=self.tensor - other.tensor,
+            )
+        else:
+            return ParityTensor(
+                edges=self.edges,
+                tensor=self.tensor - other,
+            )
+
+    def __rsub__(self, other):
+        return ParityTensor(
+            edges=self.edges,
+            tensor=other - self.tensor,
+        )
+
+    def __isub__(self, other):
+        if isinstance(other, ParityTensor):
+            self._validate_edge_compatibility(other)
+            self.tensor -= other.tensor
+        else:
+            self.tensor -= other
+        return self
+
+    def __mul__(self, other):
+        if isinstance(other, ParityTensor):  # pylint: disable=no-else-return
+            self._validate_edge_compatibility(other)
+            return ParityTensor(
+                edges=self.edges,
+                tensor=self.tensor * other.tensor,
+            )
+        else:
+            return ParityTensor(
+                edges=self.edges,
+                tensor=self.tensor * other,
+            )
+
+    def __rmul__(self, other):
+        return ParityTensor(
+            edges=self.edges,
+            tensor=other * self.tensor,
+        )
+
+    def __imul__(self, other):
+        if isinstance(other, ParityTensor):
+            self._validate_edge_compatibility(other)
+            self.tensor *= other.tensor
+        else:
+            self.tensor *= other
+        return self
+
+    def __truediv__(self, other):
+        if isinstance(other, ParityTensor):  # pylint: disable=no-else-return
+            self._validate_edge_compatibility(other)
+            return ParityTensor(
+                edges=self.edges,
+                tensor=self.tensor / other.tensor,
+            )
+        else:
+            return ParityTensor(
+                edges=self.edges,
+                tensor=self.tensor / other,
+            )
+
+    def __rtruediv__(self, other):
+        return ParityTensor(
+            edges=self.edges,
+            tensor=other / self.tensor,
+        )
+
+    def __itruediv__(self, other):
+        if isinstance(other, ParityTensor):
+            self._validate_edge_compatibility(other)
+            self.tensor /= other.tensor
+        else:
+            self.tensor /= other
+        return self
+
+    def __pos__(self):
+        return ParityTensor(
+            edges=self.edges,
+            tensor=+self.tensor,
+        )
+
+    def __neg__(self):
+        return ParityTensor(
+            edges=self.edges,
+            tensor=-self.tensor,
+        )