directly create tensor in _apply_uop (#13489)

This commit is contained in:
George Hotz
2025-11-28 19:51:06 -08:00
committed by GitHub
parent 3d76ef9ba8
commit 6f4d7c0c70
+8 -1
View File
@@ -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)