Properties

Basic

shape property

shape: tuple[sint, ...]

dtype property

dtype: DType

device property

device: str | tuple[str, ...] | None

ndim property

ndim: int

Returns the number of dimensions in the tensor.

t = Tensor([[1, 2], [3, 4]])
print(t.ndim)

2

numel

numel() -> sint

Returns the total number of elements in the tensor.

t = Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
print(t.numel())

8

Source code in tinygrad/mixin/movement.py

def numel(self) -> sint:
  """
  Returns the total number of elements in the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
  print(t.numel())
  ```
  """
  return prod(self.shape)

element_size

element_size() -> int

Returns the size in bytes of an individual element in the tensor.

t = Tensor([5], dtype=dtypes.int16)
print(t.element_size())

2

Source code in tinygrad/mixin/dtype.py

def element_size(self) -> int:
  """
  Returns the size in bytes of an individual element in the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([5], dtype=dtypes.int16)
  print(t.element_size())
  ```
  """
  if self.dtype in dtypes.weaks: raise RuntimeError(f"element_size requires a concrete dtype, got {self.dtype}")
  return self.dtype.itemsize

nbytes

nbytes() -> int

Returns the total number of bytes of all elements in the tensor.

t = Tensor([8, 9], dtype=dtypes.float)
print(t.nbytes())

8

Source code in tinygrad/mixin/op.py

def nbytes(self) -> int:
  """
  Returns the total number of bytes of all elements in the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([8, 9], dtype=dtypes.float)
  print(t.nbytes())
  ```
  """
  return int(self.numel()) * self.element_size()

is_floating_point

is_floating_point() -> bool

Returns True if the tensor contains floating point types, i.e. is one of dtypes.float64, dtypes.float32, dtypes.float16, dtypes.bfloat16.

t = Tensor([8, 9], dtype=dtypes.float32)
print(t.is_floating_point())

True

Source code in tinygrad/mixin/dtype.py

def is_floating_point(self) -> bool:
  """
  Returns `True` if the tensor contains floating point types, i.e. is one of `dtypes.float64`, `dtypes.float32`,
  `dtypes.float16`, `dtypes.bfloat16`.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([8, 9], dtype=dtypes.float32)
  print(t.is_floating_point())
  ```
  """
  return dtypes.is_float(self.dtype)

size

size(dim: int | None = None) -> sint | tuple[sint, ...]

Returns the size of the tensor. If dim is specified, return the length along dimension dim. Otherwise return the shape of the tensor.

t = Tensor([[4, 5, 6], [7, 8, 9]])
print(t.size())

(2, 3)

print(t.size(dim=1))

3

Source code in tinygrad/mixin/movement.py

def size(self, dim:int|None=None) -> sint|tuple[sint, ...]:
  """
  Returns the size of the tensor. If `dim` is specified, return the length along dimension `dim`. Otherwise return the shape of the tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[4, 5, 6], [7, 8, 9]])
  print(t.size())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.size(dim=1))
  ```
  """
  return self.shape if dim is None else self.shape[dim]

Data Access

data

data() -> memoryview

Returns the data of this tensor as a memoryview.

t = Tensor([1, 2, 3, 4])
print(np.frombuffer(t.data(), dtype=np.int32))

[1 2 3 4]

Source code in tinygrad/tensor.py

def data(self) -> memoryview:
  """
  Returns the data of this tensor as a memoryview.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 2, 3, 4])
  print(np.frombuffer(t.data(), dtype=np.int32))
  ```
  """
  if self.dtype in dtypes.weaks: return self.cast(self.commit_dtype()).data()
  if 0 in self.shape: return memoryview(bytearray(0)).cast(self.dtype.fmt)  # type: ignore[arg-type,return-value]
  assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
  buf = self._buffer()
  fmt = buf.dtype.fmt
  assert fmt is not None, f"no fmt dtype for {buf.dtype}"
  assert fmt != "e" or sys.version_info >= (3, 12)
  return buf.as_memoryview().cast(fmt, self.shape)  # type: ignore[arg-type,return-value]

item

item() -> PyConst

Returns the value of this tensor as a standard Python number.

t = Tensor(42)
print(t.item())

42

Source code in tinygrad/mixin/op.py

def item(self) -> PyConst:
  """
  Returns the value of this tensor as a standard Python number.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor(42)
  print(t.item())
  ```
  """
  assert self.numel() == 1, "must have one element for item"
  return self.data()[(0,) * len(self.shape)]

tolist

tolist() -> PyConst | list[Any]

Returns the value of this tensor as a nested list. Returns single value for const tensor.

t = Tensor([1, 2, 3, 4])
print(t.tolist())

[1, 2, 3, 4]

t = Tensor(5)
print(t.tolist())

5

Source code in tinygrad/tensor.py

def tolist(self) -> PyConst|list[Any]:
  """
  Returns the value of this tensor as a nested list.
  Returns single value for const tensor.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 2, 3, 4])
  print(t.tolist())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor(5)
  print(t.tolist())
  ```
  """
  # TODO: remove half once minimum python supports it
  if self.dtype in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s): return self.cast(dtypes.float32).tolist()
  if 0 in self.shape:
    assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
    def _tolist(shape:tuple[int, ...]): return [_tolist(shape[1:]) for _ in range(shape[0])]
    return _tolist(self.shape)
  return self.data().tolist()

numpy

numpy() -> 'numpy.ndarray'

Returns the value of this tensor as a numpy.ndarray.

t = Tensor([1, 2, 3, 4])
print(repr(t.numpy()))

array([1, 2, 3, 4], dtype=int32)

Source code in tinygrad/tensor.py

def numpy(self) -> 'numpy.ndarray':
  """
  Returns the value of this tensor as a `numpy.ndarray`.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 2, 3, 4])
  print(repr(t.numpy()))
  ```
  """
  if self.dtype in dtypes.weaks: return self.cast(self.commit_dtype()).numpy()
  assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
  import numpy as np
  if self.dtype in { dtypes.bfloat16, *dtypes.fp8s }: return self.float().numpy()
  if 0 in self.shape: return np.empty(self.shape, dtype=_to_np_dtype(self.dtype))
  return self._buffer().numpy().reshape(self.shape)

tinygrad ops

linear_with_vars

linear_with_vars(
    *lst: Tensor,
) -> tuple[UOp, dict[str, int]]

Creates the LINEAR UOp needed to realize these Tensor(s), with Variables.

Source code in tinygrad/tensor.py

@rewrite_group(lambda *tensors,ret: f"Schedule {len(tensors)} -> {len(ret[0].src)}")
def linear_with_vars(self, *lst:Tensor) -> tuple[UOp, dict[str, int]]:
  """Creates the LINEAR UOp needed to realize these Tensor(s), with Variables."""
  sink = UOp.sink(*[t.uop for t in (self,)+lst])
  if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Tensor Graph")
  # weakness ends where storage begins
  if any(u.dtype in dtypes.weaks and u.device is not None for u in sink.src):
    raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
  # The outputs and the ALLOCs beneath their wrappers get bound storage, so all aliases share storage and call dependencies.
  bases = {x.base for x in sink.src}
  for x in list(bases):
    while x.op in {Ops.STAGE, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD}: x = x.src[0].base
    if (b:=x.storage_base).op is Ops.ALLOC: bases.add(b)

  # Rebuild in dependency order: replacement values already reference the other outputs' storage.
  tensor_map:dict[UOp, UOp] = {}
  for x in sink.toposort(enter_calls=False):
    u = x.replace(src=tuple(tensor_map.get(s, s) for s in x.src))
    if x.op is Ops.ALLOC and (x.arg.bind_on_realize or x in bases): u = UOp.new_buffer(x.device, x.max_numel(), x.dtype)
    elif x in bases and u.needs_storage():
      # unwrap the rebuilt output to the compute; a STAGE means a contiguous view was requested
      src, contiguous = u, False
      while src.op in {Ops.STAGE, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD}:
        contiguous |= src.op is Ops.STAGE
        src = src.src[0]
      if src.is_virtual or src.on_disk() or 0 in src.shape or src.has_buffer_identity(after_ok=True): u = src
      elif src.op is Ops.AFTER and (src.src[1].op is Ops.STORE or (not contiguous and src.storage_base.has_buffer_identity())): u = src
      elif contiguous and (view := contiguous_mops_to_view(None, u, src)) is not None: u = view
      else:  # allocate fresh storage and store the compute into it
        buf = UOp.new_buffer(src.device, prod(src.max_shard_shape), src.dtype).reshape(src.max_shard_shape).shrink_to(src.shard_shape)
        if isinstance(src.device, tuple) and src.axis is not None: buf = buf.unshard(src.axis)
        u = buf.after(buf.store(src))
    if u is not x: tensor_map[x] = u

  sink = tensor_map.get(sink, sink)
  # Realized outputs become the storage their AFTER sequenced a store into. Compose with tensor_map before updating
  # Tensors so map values reference final storage.
  becomes_map = {u:graph_rewrite(u.src[0], pm_drop_after, bottom_up=True, name="drop after").shrink_to(u.shape)
                 for u in sink.toposort(enter_calls=False) if is_store_after(u)}
  assert not any(x in becomes_map for x in becomes_map.values())
  tensor_map = dict(zip(tensor_map, UOp.sink(*tensor_map.values()).substitute(becomes_map, walk=True).src))
  _apply_map_to_tensors(becomes_map | tensor_map, name="bufferize")

  return create_linear_with_vars(sink)

schedule_linear

schedule_linear(*lst: Tensor) -> UOp

Creates the schedule needed to realize these Tensor(s).

Source code in tinygrad/tensor.py

def schedule_linear(self, *lst:Tensor) -> UOp:
  """Creates the schedule needed to realize these Tensor(s)."""
  linear, var_vals = self.linear_with_vars(*lst)
  assert len(var_vals) == 0
  return linear

realize

realize(*lst: Tensor, do_update_stats=True) -> Tensor

Triggers the computation needed to create these Tensor(s).

Source code in tinygrad/tensor.py

@disable_gc()
def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
  """Triggers the computation needed to create these Tensor(s)."""
  to_realize = [x for x in (self,)+lst if x.uop.base.needs_storage()]
  if len(to_realize):
    linear, var_vals = Tensor.linear_with_vars(*to_realize)
    run_linear(linear, var_vals, update_stats=do_update_stats)
  return self

replace

replace(x: Tensor) -> Tensor

Replaces the data of this tensor with the data of another tensor. Only the shape of the tensors must match.

Source code in tinygrad/tensor.py

def replace(self, x:Tensor) -> Tensor:
  """
  Replaces the data of this tensor with the data of another tensor. Only the shape of the tensors must match.
  """
  # used for replacing a Tensor with a new version of it (potentially with a different device and dtype)
  assert self.shape == x.shape, f"replace shape mismatch {self.shape} != {x.shape}"
  self.uop = x.uop
  return self

assign

assign(x: Tensor | PyConst | list | tuple) -> Tensor

Source code in tinygrad/tensor.py

def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor:
  if self.dtype in dtypes.weaks: self.uop = self.uop.clone()
  is_disk = self.uop.on_disk()
  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)
  x = x._broadcast_to(self.shape)
  if x.dtype in dtypes.weaks: x = x.cast(least_upper_dtype(self.dtype, x.dtype))
  if x.dtype != self.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}")
  if isinstance(self.device, tuple) and x.uop.device is not None and self.uop.axis != x.uop.axis:
    raise RuntimeError(f"multi axis mismatch {self.uop.axis} != {x.uop.axis}")

  # TODO: this is a hack for writing to DISK. remove with working assign
  if is_disk:
    (b:=self._buffer()).copy_from(Buffer("PYTHON", b.size, b.dtype, opaque=x._data()))
    return self
  assigned_to = self.uop.storage_base
  # assigning to a value is initialization, not a write: the whole tensor is overwritten, so the pending value is dead.
  # a pending CONTIGUOUS counts only if it's the whole target: writes through views of it store into its storage
  if not assigned_to.has_buffer_identity() and (assigned_to.op is not Ops.STAGE or self.uop is assigned_to):
    self.uop = (x.uop.src[0] if x.uop.op is Ops.STAGE else x.uop).clone()
    return self
  # STORE+AFTER: STORE is the write effect (void), AFTER wraps the view for correct shape/ranging
  assign = self.uop.after(store := self.uop.store(x.uop))
  ib = self.uop
  while ib.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH} and not (ib.has_buffer_identity() and _tensor_holds(ib)): ib = ib.src[0]
  if ib is not self.uop:
    # a partial write needs storage to land in: a pending value gets explicit storage (a clone)
    target = ib if ib.has_buffer_identity(after_ok=True) else ib.clone()
    if target is not ib:
      assign = assign.substitute({ib: target}, walk=True)
      store = assign.src[1]
    # view assign: the base reads "after the store into the view" (one AFTER level). replace the node under the
    # views (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
    _apply_map_to_tensors({ib: target.after(store)}, name="Embed View Assign")
  else:
    # simple assign
    self.uop = assign
  return self

detach

detach() -> Self

Returns a new tensor with the same data as this tensor, but detached from the autograd graph.

Source code in tinygrad/mixin/elementwise.py

def detach(self) -> Self:
  """
  Returns a new tensor with the same data as this tensor, but detached from the autograd graph.
  """
  return self.alu(Ops.DETACH)

clone

clone(
    device: str | tuple[str, ...] | None = None,
) -> Tensor

Creates a clone of this tensor allocating a separate buffer for the data. If device is specified, the clone is placed on that device.

Source code in tinygrad/tensor.py

def clone(self, device:str|tuple[str, ...]|None=None) -> Tensor:
  """
  Creates a clone of this tensor allocating a separate buffer for the data.
  If `device` is specified, the clone is placed on that device.
  """
  ret = Tensor(self.uop.clone(device=device))
  if self.grad is not None: ret.grad = self.grad.clone(device=device)
  return ret.is_param_(self.is_param)

to

to(
    device: str | tuple[str, ...] | None,
    force: bool = False,
) -> Tensor

Moves the tensor to the given device. force=True inserts a transfer even for device-less values.

Source code in tinygrad/tensor.py

def to(self, device:str|tuple[str, ...]|None, force:bool=False) -> Tensor:
  """
  Moves the tensor to the given device. `force=True` inserts a transfer even for device-less values.
  """
  if self.uop.device is None and not force: return self
  if (device:=canonicalize_device(device)) == self.device: return self
  if isinstance(device, str) and is_disk_device(device):
    if isinstance(self.device, tuple): raise RuntimeError("gather to a single device before storing to DISK")
    if self.grad is not None: raise RuntimeError("tensor and gradient need separate DISK destinations; use explicit STOREs")
    dst = self.uop.empty_like(device=device)
    ret = Tensor(dst.after(dst.store(self.uop.cast(dst.dtype))))
  elif self.uop.on_creation_device(): ret = Tensor(self.uop.clone(device))
  else: ret = Tensor(self.uop.copy_to_device(device))
  if self.grad is not None: ret.grad = self.grad.to(device, force=force)
  return ret.is_param_(self.is_param)

to_

to_(device: str | tuple[str, ...] | None) -> Tensor

Moves the tensor to the given device in place.

Source code in tinygrad/tensor.py

def to_(self, device:str|tuple[str, ...]|None) -> Tensor:
  """
  Moves the tensor to the given device in place.
  """
  real = self.to(device)
  if self.grad is not None and real.grad is not None: self.grad.replace(real.grad)
  return self.replace(real)

shard

shard(
    devices: tuple[str, ...], axis: int | None = None
) -> Tensor

Shards the tensor across the given devices. Optionally specify which axis to shard on.

t = Tensor.empty(2, 4)
print(t.shard((t.device, t.device), axis=1).uop)

UOp(Ops.UNSHARD, arg=(1,), src=(
  UOp(Ops.SHRINK, arg=None, src=(
    UOp(Ops.COPY, arg=('CPU', 'CPU'), src=(
      UOp(Ops.RESHAPE, arg=None, src=(
        UOp(Ops.ALLOC, arg=ParamArg(2059, dtypes.f32, 8, device='CPU', bind_on_realize=True), src=()),
        UOp(Ops.STACK, arg=None, src=(
          x5:=UOp(Ops.CONST, arg=2, src=()),
          UOp(Ops.CONST, arg=4, src=()),)),)),
      x7:=UOp(Ops.RANGE, arg=(AxisType.DEVICE, 0), src=(
        x5,)),)),
    UOp(Ops.STACK, arg=None, src=(
      UOp(Ops.CONST, arg=0, src=()),
      UOp(Ops.MUL, arg=None, src=(
        x7,
        x5,)),)),
    UOp(Ops.STACK, arg=None, src=(
      x5,
      x5,)),)),
  x7,))

Source code in tinygrad/tensor.py

def shard(self, devices:tuple[str, ...], axis:int|None=None) -> Tensor:
  """
  Shards the tensor across the given devices. Optionally specify which axis to shard on.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.empty(2, 4)
  print(t.shard((t.device, t.device), axis=1).uop)
  ```
  """
  if self.uop.device is None: return self
  if not isinstance(self.device, str): raise RuntimeError("can't shard a multi-device tensor")
  if len(devices) == 1: return self.to(devices[0])
  devices = cast(tuple[str, ...], canonicalize_device(devices))
  # a shard of a load from a creation device (disk/npy/python) wants the copy to persist, so it inserts a clone
  src = self.uop.clone(devices) if self.uop.on_creation_device() else self.uop
  uop = src.shard(devices, None if axis is None else self._resolve_dim(axis))
  return Tensor(uop).is_param_(self.is_param)

shard_

shard_(
    devices: tuple[str, ...], axis: int | None = None
) -> Tensor

Shards the tensor across the given devices in place.

Source code in tinygrad/tensor.py

def shard_(self, devices:tuple[str, ...], axis:int|None=None) -> Tensor:
  """
  Shards the tensor across the given devices in place.
  """
  return self.replace(self.shard(devices, axis))

contiguous

contiguous() -> Self

Returns a contiguous tensor.

Source code in tinygrad/mixin/elementwise.py

def contiguous(self) -> Self:
  """
  Returns a contiguous tensor.
  """
  uop = self._uop
  if self.dtype in dtypes.weaks or uop.op is Ops.STAGE or self.device is None or uop.has_buffer_identity(): return self
  return self._wrap_uop(uop.alu(Ops.STAGE))

contiguous_backward

contiguous_backward() -> Self

Inserts a contiguous operation in the backward pass.

Source code in tinygrad/mixin/elementwise.py

def contiguous_backward(self) -> Self:
  """
  Inserts a contiguous operation in the backward pass.
  """
  return self.alu(Ops.CONTIGUOUS_BACKWARD)

Gradient

gradient

gradient(
    *targets: Self, gradient: Self | None = None
) -> list[Self]

Computes the gradient of the targets with respect to self.

x = Tensor.eye(3)
y = Tensor([[2.0,0,-2.0]])
z = y.matmul(x).sum()
dx, dy = z.gradient(x, y)

print(dx.tolist())  # dz/dx
print(dy.tolist())  # dz/dy

[[2.0, 2.0, 2.0], [0.0, 0.0, 0.0], [-2.0, -2.0, -2.0]]
[[1.0, 1.0, 1.0]]

Source code in tinygrad/mixin/op.py

def gradient(self, *targets:Self, gradient:Self|None=None) -> list[Self]:
  """
  Computes the gradient of the targets with respect to self.

  ```python exec="true" source="above" session="tensor" result="python"
  x = Tensor.eye(3)
  y = Tensor([[2.0,0,-2.0]])
  z = y.matmul(x).sum()
  dx, dy = z.gradient(x, y)

  print(dx.tolist())  # dz/dx
  print(dy.tolist())  # dz/dy
  ```
  """
  assert gradient is not None or self.shape == tuple(), "when no gradient is provided, backward must be called on a scalar tensor"
  if not (self.is_floating_point() and all(t.is_floating_point() for t in targets)): raise RuntimeError("only float Tensors have gradient")
  if any(t.dtype in dtypes.weaks for t in targets): raise RuntimeError("cannot take gradient wrt a weak Tensor")
  from tinygrad.mixin.gradient import compute_gradient
  if gradient is None: gradient = self.const_like(1.0)
  target_uops = [t._uop for t in targets]
  grads = compute_gradient(self._uop, gradient._uop, set(target_uops))
  return [self._wrap_uop(grads[x] if x in grads else x.const_like(0)) for x in target_uops]

backward

backward(gradient: Tensor | None = None) -> Tensor

Propagates the gradient of a tensor backwards through the computation graph. If the 'gradient' argument is not provided, the tensor must be a scalar, and the gradient is implicitly set to 1.0.

t = Tensor([1.0, 2.0, 3.0, 4.0])
t.sum().backward()
print(t.grad.numpy())

[1. 1. 1. 1.]

Source code in tinygrad/tensor.py

def backward(self, gradient:Tensor|None=None) -> Tensor:
  """
  Propagates the gradient of a tensor backwards through the computation graph.
  If the 'gradient' argument is not provided, the tensor must be a scalar, and the gradient is implicitly set to 1.0.
  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1.0, 2.0, 3.0, 4.0])
  t.sum().backward()
  print(t.grad.numpy())
  ```
  """
  all_uops = self.uop.toposort()
  # backward fills .grad for every in-scope float tensor with a device
  tensors_need_grad: list[Tensor] = [t for tref in all_tensors if (t:=tref()) is not None and \
                                     t.uop in all_uops and t.is_floating_point() and t.device is not None]
  # clear contexts
  for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient)):
    assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
    if g.device is None: g = g.clone(device=t.device)
    if t.grad is None: t.grad = g
    else: t.grad.assign(t.grad + g.to(t.grad.device))
  return self