mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-09-05 15:46:15 +00:00
directly create tensor in _apply_uop (#13489)
This commit is contained in:
+8
-1
@@ -173,7 +173,14 @@ class Tensor(OpMixin):
|
||||
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]
|
||||
return Tensor(new_uop, device=new_uop.device, requires_grad=True if any(needs_input_grad) else None if None in needs_input_grad else False)
|
||||
# directly create the Tensor
|
||||
ret = Tensor.__new__(Tensor)
|
||||
ret.uop = new_uop
|
||||
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
|
||||
|
||||
def _apply_broadcasted_uop(self, fxn:Callable, x:Tensor|ConstType, reverse=False) -> Tensor:
|
||||
lhs,rhs = self._broadcasted(x, reverse)
|
||||
|
||||
Reference in New Issue
Block a user