MCPcopy
hub / github.com/tinygrad/tinygrad / __setitem__

Method __setitem__

tinygrad/tensor.py:1008–1037  ·  view source on GitHub ↗
(self, indices, v:Tensor|PyConst|list|tuple)

Source from the content-addressed store, hash-verified

1006 return self._getitem(indices)
1007
1008 def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None:
1009 if isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
1010 # raise if mutation would diverge from eager (allow only pure views of a realized buffer; exclude +=/-= RHS via v_uop/v_bw)
1011 v_uop, v_bw = (v.uop, v.uop.backward_slice) if isinstance(v, Tensor) else (None, {})
1012 if self.uop.op_in_backward_slice_with_self(Ops.BUFFER):
1013 shared = self.uop.base if self.uop.base.is_realized else None
1014 if any(self.uop in t.uop.backward_slice_with_self and t.uop.base is not shared for tref in all_tensors
1015 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):
1016 raise RuntimeError("can't setitem on a tensor with other uses")
1017 if not self.uop.base.is_realized and self.is_floating_point():
1018 if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
1019 # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value
1020 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])
1021 self.replace(self._getitem(indices, v))
1022 return
1023 idx = [indices] if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)) else list(indices)
1024 is_disk = isinstance(self.device, str) and self.device.startswith("DISK")
1025 if any(isinstance(i, (Tensor, list, tuple)) for i in idx): # advanced setitem
1026 if is_disk: raise RuntimeError("advanced setitem is not supported for DISK tensors")
1027 if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
1028 self.assign(self._getitem(indices, v))
1029 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
1030 view = self[indices]
1031 if isinstance(v, Tensor) and v.uop.op is Ops.AFTER and v.uop in view.uop.base.src: return
1032 view.assign(v)
1033 else: # basic setitem, self is not realized
1034 if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
1035 # __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value
1036 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])
1037 self.replace(self._getitem(indices, v))
1038
1039 def __delitem__(self, indices) -> None:
1040 raise TypeError("Tensor does not support deleting items")

Callers

nothing calls this directly

Calls 9

_apply_uopMethod · 0.95
replaceMethod · 0.95
_getitemMethod · 0.95
assignMethod · 0.95
all_intFunction · 0.90
TensorClass · 0.85
is_floating_pointMethod · 0.80

Tested by

no test coverage detected