add magic slot for dtensor

Signed-off-by: weiwee <wbwmat@gmail.com>
FederatedAI · Jun 19, 2023 · f4ea541 · f4ea541
1 parent 3552824
commit f4ea541
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 100 deletions.
diff --git a/python/fate/arch/context/io/model/__init__.py b/python/fate/arch/context/io/model/__init__.py
diff --git a/python/fate/arch/context/io/model/file.py b/python/fate/arch/context/io/model/file.py
diff --git a/python/fate/arch/context/io/model/http.py b/python/fate/arch/context/io/model/http.py
diff --git a/python/fate/arch/tensor/distributed/_tensor.py b/python/fate/arch/tensor/distributed/_tensor.py
@@ -47,6 +47,36 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
     def __init__(self, shardings: "Shardings") -> None:
         self.shardings = shardings
 
+    def __add__(self, other):
+        return torch.add(self, other)
+
+    def __radd__(self, other):
+        return torch.add(other, self)
+
+    def __sub__(self, other):
+        return torch.sub(self, other)
+
+    def __rsub__(self, other):
+        return torch.rsub(self, other)
+
+    def __mul__(self, other):
+        return torch.mul(self, other)
+
+    def __rmul__(self, other):
+        return torch.mul(other, self)
+
+    def __truediv__(self, other):
+        return torch.div(self, other)
+
+    def __rtruediv__(self, other):
+        return torch.div(other, self)
+
+    def __matmul__(self, other):
+        return torch.matmul(self, other)
+
+    def __rmatmul__(self, other):
+        return torch.matmul(other, self)
+
     @property
     def shape(self):
         return self.shardings.shape