Skip to content

Properties

Basic¤

shape property ¤

shape: Tuple[sint, ...]

dtype property ¤

dtype: DType

device property ¤

device: Union[str, Tuple[str, ...]]

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/tensor.py
3515
3516
3517
3518
3519
3520
3521
3522
3523
3524
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/tensor.py
3526
3527
3528
3529
3530
3531
3532
3533
3534
3535
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())
  ```
  """
  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/tensor.py
3537
3538
3539
3540
3541
3542
3543
3544
3545
3546
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 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 dtype.float64, dtype.float32, dtype.float16, dtype.bfloat16.

t = Tensor([8, 9], dtype=dtypes.float32)
print(t.is_floating_point())
True
Source code in tinygrad/tensor.py
3548
3549
3550
3551
3552
3553
3554
3555
3556
3557
3558
def is_floating_point(self) -> bool:
  """
  Returns `True` if the tensor contains floating point types, i.e. is one of `dtype.float64`, `dtype.float32`,
  `dtype.float16`, `dtype.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: Optional[int] = None,
) -> Union[sint, Tuple[sint, ...]]

Return 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/tensor.py
3560
3561
3562
3563
3564
3565
3566
3567
3568
3569
3570
3571
3572
def size(self, dim:Optional[int]=None) -> Union[sint, Tuple[sint, ...]]:
  """
  Return 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
272
273
274
275
276
277
278
279
280
281
282
283
284
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))
  ```
  """
  assert self.dtype.base.fmt is not None, f"no fmt dtype for {self.dtype.base}"
  assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
  if TYPE_CHECKING or sys.version_info < (3, 12): assert self.dtype.base.fmt != "e"
  return self._data().cast(self.dtype.base.fmt) if 0 in self.shape else self._data().cast(self.dtype.base.fmt, self.shape)

item ¤

item() -> ConstType

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

t = Tensor(42)
print(t.item())
42
Source code in tinygrad/tensor.py
286
287
288
289
290
291
292
293
294
295
296
def item(self) -> ConstType:
  """
  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 ¤

Returns the value of this tensor as a nested list.

t = Tensor([1, 2, 3, 4])
print(t.tolist())
[1, 2, 3, 4]
Source code in tinygrad/tensor.py
300
301
302
303
304
305
306
307
308
309
def tolist(self) -> Union[Sequence[ConstType], ConstType]:
  """
  Returns the value of this tensor as a nested list.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 2, 3, 4])
  print(t.tolist())
  ```
  """
  return self.data().tolist()

numpy ¤

numpy() -> 'np.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
311
312
313
314
315
316
317
318
319
320
321
322
323
324
def numpy(self) -> 'np.ndarray':  # type: ignore [name-defined] # noqa: F821
  """
  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()))
  ```
  """
  import numpy as np
  if self.dtype.base == dtypes.bfloat16: return self.float().numpy()
  assert _to_np_dtype(self.dtype.base) is not None, f"no np dtype for {self.dtype.base}"
  assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
  return np.frombuffer(self._data(), dtype=_to_np_dtype(self.dtype.base)).reshape(self.shape)

tinygrad ops¤

schedule_with_vars ¤

schedule_with_vars(
    *lst: Tensor,
) -> Tuple[List[ScheduleItem], Dict[Variable, int]]

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

Note

A Tensor can only be scheduled once.

Source code in tinygrad/tensor.py
209
210
211
212
213
214
215
216
def schedule_with_vars(self, *lst:Tensor) -> Tuple[List[ScheduleItem], Dict[Variable, int]]:
  """
  Creates the schedule needed to realize these Tensor(s), with Variables.

  NOTE: A Tensor can only be scheduled once.
  """
  schedule, var_vals = create_schedule_with_vars(flatten([x.lazydata.lbs for x in (self,)+lst]))
  return memory_planner(schedule), var_vals

schedule ¤

schedule(*lst: Tensor) -> List[ScheduleItem]

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

Source code in tinygrad/tensor.py
218
219
220
221
222
def schedule(self, *lst:Tensor) -> List[ScheduleItem]:
  """Creates the schedule needed to realize these Tensor(s)."""
  schedule, var_vals = self.schedule_with_vars(*lst)
  assert len(var_vals) == 0
  return schedule

realize ¤

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

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

Source code in tinygrad/tensor.py
224
225
226
227
def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
  """Triggers the computation needed to create these Tensor(s)."""
  run_schedule(*self.schedule_with_vars(*lst), do_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
229
230
231
232
233
234
235
236
237
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 not x.requires_grad and getattr(self, '_ctx', None) is None
  assert self.shape == x.shape, f"replace shape mismatch {self.shape} != {x.shape}"
  self.lazydata = x.lazydata
  return self

assign ¤

assign(x) -> Tensor
Source code in tinygrad/tensor.py
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
def assign(self, x) -> Tensor:
  # TODO: this is a hack for writing to DISK. remove with working assign
  if isinstance(self.device, str) and self.device.startswith("DISK"):
    if x.__class__ is not Tensor: x = Tensor(x, device="CLANG", dtype=self.dtype)
    self.contiguous().realize().lazydata.base.realized.copyin(x._data())
    return self
  if x.__class__ is not Tensor: x = Tensor(x, device=self.device, dtype=self.dtype)
  if DEBUG >= 4: print(f"assign {self.lazydata} <- {x.lazydata}")
  if self.lazydata is x.lazydata: return self  # a self assign is a NOOP
  # NOTE: we allow cross device assign
  assert self.shape == x.shape, f"assign shape mismatch {self.shape} != {x.shape}"
  assert self.device == x.device, f"assign device mismatch {self.device} != {x.device}"
  assert self.dtype == x.dtype, f"assign dtype mismatch {self.dtype} != {x.dtype}"
  assert not isinstance(self.lazydata, MultiLazyBuffer) or self.lazydata.axis == x.lazydata.axis, "axis must match on MultiLazyBuffer"
  assert not x.requires_grad  # self requires_grad is okay?
  if not self.lazydata.is_realized: return self.replace(x)
  self.lazydata = self.lazydata.assign(x.lazydata)
  return self

detach ¤

detach() -> Tensor

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

Source code in tinygrad/tensor.py
258
259
260
261
262
def detach(self) -> Tensor:
  """
  Returns a new tensor with the same data as this tensor, but detached from the autograd graph.
  """
  return Tensor(self.lazydata, device=self.device, requires_grad=False)

to ¤

to(device: Optional[Union[str, Tuple[str, ...]]]) -> Tensor

Moves the tensor to the given device.

Source code in tinygrad/tensor.py
335
336
337
338
339
340
341
342
343
344
345
def to(self, device:Optional[Union[str, Tuple[str, ...]]]) -> Tensor:
  """
  Moves the tensor to the given device.
  """
  device = tuple(Device.canonicalize(x) for x in device) if isinstance(device, (tuple, list)) else Device.canonicalize(device)
  if device == self.device: return self
  if not isinstance(device, str): return self.shard(device)
  ret = Tensor(self.lazydata, device, requires_grad=self.requires_grad)
  if self.grad is not None: ret.grad = self.grad.to(device)
  if hasattr(self, '_ctx'): ret._ctx = self._ctx
  return ret

to_ ¤

to_(device: Optional[Union[str, Tuple[str, ...]]])

Moves the tensor to the given device in place.

Source code in tinygrad/tensor.py
347
348
349
350
351
352
353
354
def to_(self, device:Optional[Union[str, Tuple[str, ...]]]):
  """
  Moves the tensor to the given device in place.
  """
  real = self.to(device)
  # TODO: is this assign?
  if self.grad is not None and real.grad is not None: self.grad.lazydata = real.grad.lazydata
  self.lazydata = real.lazydata

shard ¤

shard(
    devices: Tuple[str, ...],
    axis: Optional[int] = None,
    splits: Optional[Tuple[int, ...]] = None,
) -> Tensor

Shards the tensor across the given devices. Optionally specify which axis to shard on, and how to split it across devices.

t = Tensor.empty(2, 3)
print(t.shard((t.device, t.device), axis=1, splits=(2, 1)).lazydata)
<MLB self.axis=1 self.real=[True, True] 
CLANG ShapeTracker(views=(View(shape=(2, 2), strides=(2, 1), offset=0, mask=None, contiguous=True),))
CLANG ShapeTracker(views=(View(shape=(2, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),))>
Source code in tinygrad/tensor.py
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
def shard(self, devices:Tuple[str, ...], axis:Optional[int]=None, splits:Optional[Tuple[int, ...]]=None) -> Tensor:
  """
  Shards the tensor across the given devices. Optionally specify which axis to shard on, and how to split it across devices.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.empty(2, 3)
  print(t.shard((t.device, t.device), axis=1, splits=(2, 1)).lazydata)
  ```

  """
  assert isinstance(self.lazydata, LazyBuffer), "can't shard a MultiLazyBuffer"
  devices, bounds = tuple(Device.canonicalize(x) for x in devices), None
  if axis is not None:
    if axis < 0: axis += len(self.shape)
    if splits is None:
      if not isinstance(total:=self.shape[axis], int): raise RuntimeError(f"cannot shard symbolic shape {self.shape=}, {axis=}")
      sz = ceildiv(total, len(devices))
      splits = tuple([max(0, min(sz, total - sz*i)) for i in range(len(devices))])
    assert sum(splits) == self.shape[axis], "specified splits do not sum up to axis shape"
    boundaries = tuple(itertools.accumulate(splits))
    bounds = tuple(zip((0,) + boundaries, boundaries))
  return Tensor(MultiLazyBuffer.from_sharded(self.lazydata, devices, axis, bounds), device=devices, requires_grad=self.requires_grad)

shard_ ¤

shard_(
    devices: Tuple[str, ...],
    axis: Optional[int] = None,
    splits: Optional[Tuple[int, ...]] = None,
)

Shards the tensor across the given devices in place.

Source code in tinygrad/tensor.py
379
380
381
382
383
384
def shard_(self, devices:Tuple[str, ...], axis:Optional[int]=None, splits:Optional[Tuple[int, ...]]=None):
  """
  Shards the tensor across the given devices in place.
  """
  self.lazydata = self.shard(devices, axis, splits).lazydata
  return self

contiguous ¤

contiguous()

Returns a contiguous tensor.

Source code in tinygrad/tensor.py
2368
2369
2370
2371
2372
def contiguous(self):
  """
  Returns a contiguous tensor.
  """
  return F.Contiguous.apply(self)

contiguous_backward ¤

contiguous_backward()

Inserts a contiguous operation in the backward pass.

Source code in tinygrad/tensor.py
2373
2374
2375
2376
2377
def contiguous_backward(self):
  """
  Inserts a contiguous operation in the backward pass.
  """
  return F.ContiguousBackward.apply(self)

backward ¤

backward(
    gradient: Optional[Tensor] = None,
    retain_graph: bool = False,
) -> 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. If 'retain_graph' is false, the graph used to compute the grads will be freed. Otherwise, it will be kept. Keeping it can increase memory usage.

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

Source code in tinygrad/tensor.py
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
def backward(self, gradient:Optional[Tensor]=None, retain_graph:bool=False) -> 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.
  If 'retain_graph' is false, the graph used to compute the grads will be freed. Otherwise, it will be kept. Keeping it can increase memory usage.
  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
  t.sum().backward()
  print(t.grad.numpy())
  ```
  """
  toposorted = self._deepwalk()
  if gradient is None:
    assert self.shape == tuple(), "when no gradient is provided, backward must be called on a scalar tensor"
    # fill in the first grad with one. don't use Tensor.ones because we don't need contiguous
    # this is "implicit gradient creation"
    gradient = Tensor(1.0, dtype=self.dtype, device=self.device, requires_grad=False)

  assert self.shape == gradient.shape, f"grad shape must match tensor shape, {gradient.shape!r} != {self.shape!r}"
  self.grad = gradient
  for t0 in reversed(toposorted):
    if t0.grad is None: raise RuntimeError(f"tensor {t0} has no grad")
    token = _METADATA.set(dataclasses.replace(md, backward=True) if (md := t0._ctx.metadata) is not None else None)
    grads = t0._ctx.backward(t0.grad.lazydata)
    _METADATA.reset(token)
    grads = [Tensor(g, device=self.device, requires_grad=False) if g is not None else None
      for g in ([grads] if len(t0._ctx.parents) == 1 else grads)]
    for t, g in zip(t0._ctx.parents, grads):
      if g is not None and t.requires_grad:
        assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
        t.grad = g if t.grad is None else (t.grad + g)
    if not retain_graph: del t0._ctx
  return self