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
3332
3333
3334
3335
3336
3337
3338
3339
3340
3341
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
3343
3344
3345
3346
3347
3348
3349
3350
3351
3352
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
3354
3355
3356
3357
3358
3359
3360
3361
3362
3363
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
3365
3366
3367
3368
3369
3370
3371
3372
3373
3374
3375
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
3377
3378
3379
3380
3381
3382
3383
3384
3385
3386
3387
3388
3389
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
269
270
271
272
273
274
275
276
277
278
279
280
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.fmt is not None, f"no fmt dtype for {self.dtype}"
  assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
  return self._data().cast(self.dtype.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
282
283
284
285
286
287
288
289
290
291
292
293
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.dtype.fmt is not None, f"no fmt dtype for {self.dtype}"
  assert self.numel() == 1, "must have one element for item"
  return self._data().cast(self.dtype.fmt)[0]

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
297
298
299
300
301
302
303
304
305
306
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
308
309
310
311
312
313
314
315
316
317
318
319
320
321
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 == dtypes.bfloat16: return self.float().numpy()
  assert _to_np_dtype(self.dtype) is not None, f"no np dtype for {self.dtype}"
  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)).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
206
207
208
209
210
211
212
213
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
215
216
217
218
219
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
221
222
223
224
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
226
227
228
229
230
231
232
233
234
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
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
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="NPY", dtype=self.dtype)
    self.contiguous().realize().lazydata.base.realized.copyin(x.numpy().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
255
256
257
258
259
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
332
333
334
335
336
337
338
339
340
341
342
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
344
345
346
347
348
349
350
351
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
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
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 = round_up(total, len(devices)) // 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
376
377
378
379
380
381
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
2272
2273
2274
2275
2276
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
2277
2278
2279
2280
2281
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
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
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