tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-31 00:26:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

minor _apply_uop cleanup (#14180)

Browse Source 
give fxn a return type and minor style change
This commit is contained in:
chenyu
2026-01-16 11:27:55 -05:00
committed by GitHub
parent e855ec8ee3
commit 6790165ef8
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+6 -6
View File
@@ -175,15 +175,15 @@ class Tensor(OpMixin):
@suppress_finalizing
def __del__(self): all_tensors.pop(weakref.ref(self), None)
def _apply_uop(self, fxn:Callable, *x:Tensor, extra_args=(), **kwargs) -> Tensor:
new_uop: UOp = fxn(*[t.uop for t in (self,)+x], *extra_args, **kwargs)
if (metadata:=_METADATA.get()) is not None and TRACEMETA >= 1: all_metadata[new_uop] = (metadata,)
needs_input_grad = [t.requires_grad for t in (self,)+x]
def _apply_uop(self, fxn:Callable[..., UOp], *x:Tensor, extra_args=(), **kwargs) -> Tensor:
srcs = (self,)+x
new_uop: UOp = fxn(*[t.uop for t in srcs], *extra_args, **kwargs)
if TRACEMETA >= 1 and (metadata:=_METADATA.get()) is not None: all_metadata[new_uop] = (metadata,)
needs_input_grad = [t.requires_grad for t in srcs]
# directly create the Tensor
ret = Tensor.__new__(Tensor)
ret.uop = new_uop
ret.uop, ret.grad = new_uop, None
ret.requires_grad = True if any(needs_input_grad) else None if None in needs_input_grad else False
ret.grad = None
# add to all_tensors after construction succeeds
all_tensors[weakref.ref(ret)] = None
return ret