github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

Tensor(uop) does not need explicit device (#15361)

Browse Source
This commit is contained in:
chenyu
2026-03-19 00:44:33 -04:00
committed by GitHub
parent f32c2e43a7
commit 92dfef8060
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tinygrad/tensor.py
View File

@@ -238,10 +238,10 @@ class Tensor(OpMixin):
param = UOp.param(slot, self.dtype, self.uop.shard_shape, self.device).multi(self.uop.axis)
else:
param = UOp.param(slot, self.dtype, self.shape, self.device)
return Tensor(param, device=self.device)
return Tensor(param)
def call(self, *lst:Tensor, fxn:Tensor|UOp, grad_fxn:Callable|None=None) -> Tensor:
fret = (fxn.uop if isinstance(fxn, Tensor) else fxn).call(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn)
return Tensor(fret.gettuple(0), device=self.device)
return Tensor(fret.gettuple(0))
def custom_kernel(self, *lst:Tensor, fxn:Callable, grad_fxn:Callable|None=None) -> list[Tensor]:
"""
@@ -324,7 +324,7 @@ class Tensor(OpMixin):
"""
Returns a new tensor with the same data as this tensor, but detached from the autograd graph.
"""
return Tensor(self.uop.detach(), device=self.device, requires_grad=False)
return Tensor(self.uop.detach(), requires_grad=False)
def _buffer(self) -> Buffer:
from tinygrad.engine.realize import capturing
@@ -803,7 +803,7 @@ class Tensor(OpMixin):
if kwargs.get("device") is not None: raise RuntimeError("cannot specify `device` on `*_like` of a multi device tensor")
if self.uop.axis is None: return fxn(self.shape, *args, dtype=dtype, **kwargs).shard(self.device)
stacked = UOp(Ops.MSTACK, dtype=dtype, src=tuple([fxn(self.uop.shard_shape, *args, device=d, dtype=dtype, **kwargs).uop for d in self.device]))
return Tensor(UOp.multi(stacked, axis=self.uop.axis), device=self.device, dtype=dtype)
return Tensor(UOp.multi(stacked, axis=self.uop.axis), dtype=dtype)
def full_like(self, fill_value:PyConst, **kwargs) -> Tensor:
"""