Commit 0c0489c5 authored by pengcheng888's avatar pengcheng888
Browse files

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

parent cad2d45a
......@@ -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 ''})"
......
......@@ -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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment