Merge pull request #681 from pengcheng888/issue/679_1

issue/679 - 减少创建和使用infini.tensor对象过程中的属性调用次数，以降低耗时

python/infinicore/device.py

@@ -2,16 +2,20 @@ from infinicore.lib import _infinicore
 
 
 class device:
-    def __init__(self, type=None, index=None):
-        if type is None:
-            type = "cpu"
+    # Public attributes describing the device
+    type: str
+    index: int
+    _underlying: _infinicore.Device
 
+    def __init__(self, type=None, index=None):
         if isinstance(type, device):
             self.type = type.type
             self.index = type.index
-
             return
 
+        if type is None:
+            type = "cpu"
+
         if ":" in type:
             if index is not None:
                 raise ValueError(
@@ -22,12 +26,14 @@ class device:
             index = int(index)
 
         self.type = type
-
-        self.index = index
-
-        _type, _index = device._to_infinicore_device(type, index if index else 0)
-
-        self._underlying = _infinicore.Device(_type, _index)
+        self.index = index if index else 0
+
+    def __getattr__(self, name):
+        # Lazily construct and cache an attribute.
+        # such as, self._underlying .
+        _type, _index = device._to_infinicore_device(self.type, self.index)
+        setattr(self, name, _infinicore.Device(_type, _index))
+        return getattr(self, name)
 
     def __repr__(self):
         return f"device(type='{self.type}'{f', index={self.index}' if self.index is not None else ''})"

python/infinicore/tensor.py

@@ -14,30 +14,35 @@ from .utils import (
 
 
 class Tensor:
+    # Public attributes describing the device
+    _underlying: _infinicore.Tensor
+    _torch_ref: "torch.Tensor"  # noqa: F821
+    shape: list[int]
+    dtype: infinicore.dtype
+    device: infinicore.device
+
     def __init__(self, underlying, *, _torch_ref=None):
         """An internal method. Please do not use this directly."""
 
         self._underlying = underlying
-
-        self._dtype = infinicore.dtype(self._underlying.dtype)
-
-        self._device = infinicore.device._from_infinicore_device(
-            self._underlying.device
-        )
-
         self._torch_ref = _torch_ref
 
-    @property
-    def shape(self):
-        return self._underlying.shape
-
-    @property
-    def dtype(self):
-        return self._dtype
-
-    @property
-    def device(self):
-        return self._device
+    def __getattr__(self, name):
+        # Lazily construct and cache an attribute.
+        # such as, self.shape, self.dtype, self.device .
+        if name == "shape":
+            setattr(self, name, getattr(self._underlying, name))
+        elif name == "dtype":
+            setattr(self, name, infinicore.dtype(getattr(self._underlying, name)))
+        elif name == "device":
+            setattr(
+                self,
+                name,
+                infinicore.device._from_infinicore_device(
+                    getattr(self._underlying, name)
+                ),
+            )
+        return getattr(self, name)
 
     @property
     def ndim(self):
@@ -101,6 +106,10 @@ class Tensor:
     def __add__(self, other):
         return infinicore.add(self, other)
 
+    def __iadd__(self, other):
+        infinicore.add(self, other, out=self)
+        return self
+
     def __matmul__(self, other):
         return infinicore.matmul(self, other)