From 3ae0e282a7087ceab3b77261484eef6406ccc75a Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Mon, 11 Aug 2025 01:56:38 +0800 Subject: [PATCH] Add check for type error for inplace arithmetic operators. --- grassmann_tensor/tensor.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/grassmann_tensor/tensor.py b/grassmann_tensor/tensor.py index c1bd774..d8bf085 100644 --- a/grassmann_tensor/tensor.py +++ b/grassmann_tensor/tensor.py @@ -391,8 +391,11 @@ def __iadd__(self, other: typing.Any) -> GrassmannTensor: if isinstance(other, GrassmannTensor): self._validate_edge_compatibility(other) self._tensor += other._tensor - else: + return self + try: self._tensor += other + except TypeError: + return NotImplemented if isinstance(self._tensor, torch.Tensor): return self return NotImplemented @@ -431,8 +434,11 @@ def __isub__(self, other: typing.Any) -> GrassmannTensor: if isinstance(other, GrassmannTensor): self._validate_edge_compatibility(other) self._tensor -= other._tensor - else: + return self + try: self._tensor -= other + except TypeError: + return NotImplemented if isinstance(self._tensor, torch.Tensor): return self return NotImplemented @@ -471,8 +477,11 @@ def __imul__(self, other: typing.Any) -> GrassmannTensor: if isinstance(other, GrassmannTensor): self._validate_edge_compatibility(other) self._tensor *= other._tensor - else: + return self + try: self._tensor *= other + except TypeError: + return NotImplemented if isinstance(self._tensor, torch.Tensor): return self return NotImplemented @@ -511,8 +520,11 @@ def __itruediv__(self, other: typing.Any) -> GrassmannTensor: if isinstance(other, GrassmannTensor): self._validate_edge_compatibility(other) self._tensor /= other._tensor - else: + return self + try: self._tensor /= other + except TypeError: + return NotImplemented if isinstance(self._tensor, torch.Tensor): return self return NotImplemented