diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index be663c59a4..906147c4fa 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -404,10 +404,6 @@ class OpMixin(ElementwiseMixin, ReduceMixin): if x.dtype == y.dtype or any(isinstance(d, PtrDType) for d in (x.dtype, y.dtype)): return x, y return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype) - def _binop(self, op:Ops, x, reverse:bool) -> Self: - lhs, rhs = self._broadcasted(x, reverse) - return lhs.alu(op, rhs) - def dot(self, w:Self, dtype:DTypeLike|None=None) -> Self: """ Performs dot product between two tensors. @@ -1304,9 +1300,9 @@ class OpMixin(ElementwiseMixin, ReduceMixin): ``` """ axis = tuple(range(-len(k_ := make_tuple(kernel_size, 2)), 0)) - pads = resolve_pool_pads(padding, len(k_)) - if ceil_mode: pads = self._apply_ceil_mode(pads, k_, stride if stride is not None else k_, dilation) s_ = stride if stride is not None else k_ + pads = resolve_pool_pads(padding, len(k_)) + if ceil_mode: pads = self._apply_ceil_mode(pads, k_, s_, dilation) pooled = self._pad_constant(((0,0),)*(self.ndim-len(k_)) + flat_to_grouped(pads), self.dtype.min)._pool(k_, s_, dilation) if not return_indices: return pooled.max(axis) spatial_sz = int(prod(spatial_shape := self.shape[-len(k_):])) diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index 0a4fd08779..21a1ee4215 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -20,7 +20,8 @@ class ElementwiseMixin(DTypeMixin, CreationMixin): return self.const_like(x) if not isinstance(x, ElementwiseMixin) else x def _binop(self, op: Ops, x: Self | ConstType, reverse: bool) -> Self: - return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x)) + lhs, rhs = self._broadcasted(x, reverse) + return lhs.alu(op, rhs) def usum(self, *uops) -> Self: return functools.reduce(operator.or_ if self.dtype is dtypes.bool else operator.add, argfix(*uops), self) def uprod(self, *uops) -> Self: return functools.reduce(operator.and_ if self.dtype is dtypes.bool else operator.mul, argfix(*uops), self) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index d7c66b4ae4..027e2064b2 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -241,8 +241,7 @@ class Tensor(RandMixin): if not isinstance(x, Tensor): x = Tensor(x, device="CPU" if is_disk else self.device, dtype=self.dtype) if self.uop is x.uop: return self # a self assign is a NOOP # broadcast x (shape only, dtype must match) - if self.shape != x.shape: x = x._broadcast_to(self.shape) - if self.shape != x.shape: raise RuntimeError(f"assign shape mismatch {self.shape} != {x.shape}") + x = x._broadcast_to(self.shape) if not is_disk and x.uop.device is not None and self.device is not None and self.device != x.device: raise RuntimeError(f"assign device mismatch {self.device} != {x.device}") if not is_disk and self.dtype != x.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}") @@ -469,7 +468,7 @@ class Tensor(RandMixin): Creates an empty tensor with the same shape as `self`. If `dtype` is not specified, the dtype of `self` is used. """ - return Tensor(self.uop.empty_like(dtype, self.device if device is None else device), **kwargs) + return Tensor(self.uop.empty_like(dtype, device), **kwargs) @staticmethod def from_blob(ptr:int, shape:tuple[int, ...], **kwargs) -> Tensor: @@ -866,27 +865,23 @@ class Tensor(RandMixin): if any(self.uop in t.uop.backward_slice_with_self and t.uop.base is not shared for tref in all_tensors if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw): raise RuntimeError("can't setitem on a tensor with other uses") - if not self.uop.base.is_realized and self.is_floating_point(): + idx = [indices] if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)) else list(indices) + is_disk = isinstance(self.device, str) and self.device.startswith("DISK") + advanced = any(isinstance(i, (Tensor, list, tuple)) for i in idx) + realized = is_disk or self.uop.is_realized or self.uop.base.op is Ops.BUFFER or self.uop._base_buffer_is_realized() + if (not self.uop.base.is_realized and self.is_floating_point()) or not (advanced or realized): if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype) # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value if v.uop.op is Ops.AFTER and any(s.op is Ops.STORE for s in v.uop.src[1:]): v = v._apply_uop(lambda x: x.src[1].src[1]) self.replace(self._getitem(indices, v)) - return - idx = [indices] if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)) else list(indices) - is_disk = isinstance(self.device, str) and self.device.startswith("DISK") - if any(isinstance(i, (Tensor, list, tuple)) for i in idx): # advanced setitem + elif advanced: # advanced setitem if is_disk: raise RuntimeError("advanced setitem is not supported for DISK tensors") if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype) self.assign(self._getitem(indices, v)) - elif is_disk or self.uop.is_realized or self.uop.base.op is Ops.BUFFER or self.uop._base_buffer_is_realized(): # basic setitem + else: # basic setitem view = self[indices] if isinstance(v, Tensor) and v.uop.op is Ops.AFTER and v.uop in view.uop.base.src: return view.assign(v) - else: # basic setitem, self is not realized - if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype) - # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value - if v.uop.op is Ops.AFTER and any(s.op is Ops.STORE for s in v.uop.src[1:]): v = v._apply_uop(lambda x: x.src[1].src[1]) - self.replace(self._getitem(indices, v)) def __delitem__(self, indices) -> None: raise TypeError("Tensor does not support deleting items")