Complex Ops

Reduce

sum

sum(
    axis: int | Sequence[int] | None = None,
    keepdim=False,
    dtype: DTypeLike | None = None,
) -> Self

Returns the sum of the elements of the tensor along the specified axis or axes.

You can pass in axis and keepdim keyword arguments to control the axis along which the maximum is computed and whether the reduced dimensions are retained.

You can pass in dtype keyword argument to control the data type of the accumulation. If not specified, the accumulation data type is chosen based on the input tensor's data type.

t = Tensor.arange(6).reshape(2, 3)
print(t.numpy())

[[0 1 2]
 [3 4 5]]

print(t.sum().numpy())

15

print(t.sum(axis=0).numpy())

[3 5 7]

print(t.sum(axis=1).numpy())

[ 3 12]

Source code in tinygrad/mixin/reduce.py

def sum(self, axis:int|Sequence[int]|None=None, keepdim=False, dtype:DTypeLike|None=None) -> Self:
  """
  Returns the sum of the elements of the tensor along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the maximum is computed and whether the reduced dimensions are retained.

  You can pass in `dtype` keyword argument to control the data type of the accumulation.
  If not specified, the accumulation data type is chosen based on the input tensor's data type.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(6).reshape(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.sum().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.sum(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.sum(axis=1).numpy())
  ```
  """
  ret = self.cast(sum_acc_dtype(self.dtype) if dtype is None else to_dtype(dtype))._reduce(Ops.ADD, axis, keepdim)
  return ret.cast(self.dtype) if dtype is None and self.dtype in (dtypes.float16, dtypes.bfloat16, *dtypes.fp8s) else ret

prod

prod(
    axis: int | Sequence[int] | None = None,
    keepdim=False,
    dtype: DTypeLike | None = None,
) -> Self

Returns the product of the elements of the tensor along the specified axis or axes.

You can pass in axis and keepdim keyword arguments to control the axis along which the maximum is computed and whether the reduced dimensions are retained.

You can pass in dtype keyword argument to control the data type of the accumulation. If not specified, the accumulation data type is chosen based on the input tensor's data type.

t = Tensor([-1, -2, -3, 1, 2, 3]).reshape(2, 3)
print(t.numpy())

[[-1 -2 -3]
 [ 1  2  3]]

print(t.prod().numpy())

-36

print(t.prod(axis=0).numpy())

[-1 -4 -9]

print(t.prod(axis=1).numpy())

[-6  6]

Source code in tinygrad/mixin/reduce.py

def prod(self, axis:int|Sequence[int]|None=None, keepdim=False, dtype:DTypeLike|None=None) -> Self:
  """
  Returns the product of the elements of the tensor along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the maximum is computed and whether the reduced dimensions are retained.

  You can pass in `dtype` keyword argument to control the data type of the accumulation.
  If not specified, the accumulation data type is chosen based on the input tensor's data type.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([-1, -2, -3, 1, 2, 3]).reshape(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.prod().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.prod(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.prod(axis=1).numpy())
  ```
  """
  return self.cast(to_dtype(dtype) if dtype is not None else self.dtype)._reduce(Ops.MUL, axis, keepdim)

max

max(
    axis: int | Sequence[int] | None = None, keepdim=False
) -> Self

Returns the maximum value of the tensor along the specified axis or axes.

You can pass in axis and keepdim keyword arguments to control the axis along which the maximum is computed and whether the reduced dimensions are retained.

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

[[1 0 2]
 [5 4 3]]

print(t.max().numpy())

5

print(t.max(axis=0).numpy())

[5 4 3]

print(t.max(axis=1, keepdim=True).numpy())

[[2]
 [5]]

Source code in tinygrad/mixin/reduce.py

def max(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Self:
  """
  Returns the maximum value of the tensor along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the maximum is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 0, 2], [5, 4, 3]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.max().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.max(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.max(axis=1, keepdim=True).numpy())
  ```
  """
  return self._reduce(Ops.MAX, axis, keepdim)

min

min(
    axis: int | Sequence[int] | None = None, keepdim=False
) -> Tensor

Returns the minimum value of the tensor along the specified axis or axes.

You can pass in axis and keepdim keyword arguments to control the axis along which the minimum is computed and whether the reduced dimensions are retained.

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

[[1 0 2]
 [5 4 3]]

print(t.min().numpy())

0

print(t.min(axis=0).numpy())

[1 0 2]

print(t.min(axis=1, keepdim=True).numpy())

[[0]
 [3]]

Source code in tinygrad/tensor.py

def min(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
  """
  Returns the minimum value of the tensor along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the minimum is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 0, 2], [5, 4, 3]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.min().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.min(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.min(axis=1, keepdim=True).numpy())
  ```
  """
  return self._inverse().max(axis=axis, keepdim=keepdim)._inverse()

any

any(
    axis: int | Sequence[int] | None = None, keepdim=False
) -> Tensor

Tests if any element evaluates to True along the specified axis or axes.

You can pass in axis and keepdim keyword arguments to control the reduce axis and whether the reduced dimensions are retained.

t = Tensor([[True, True], [True, False], [False, False]])
print(t.numpy())

[[ True  True]
 [ True False]
 [False False]]

print(t.any().numpy())

True

print(t.any(axis=0).numpy())

[ True  True]

print(t.any(axis=1, keepdim=True).numpy())

[[ True]
 [ True]
 [False]]

Source code in tinygrad/tensor.py

def any(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
  """
  Tests if any element evaluates to `True` along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the reduce axis and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[True, True], [True, False], [False, False]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.any().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.any(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.any(axis=1, keepdim=True).numpy())
  ```
  """
  return self.bool().max(axis, keepdim)

all

all(
    axis: int | Sequence[int] | None = None, keepdim=False
) -> Tensor

Tests if all element evaluates to True along the specified axis or axes.

You can pass in axis and keepdim keyword arguments to control the reduce axis and whether the reduced dimensions are retained.

t = Tensor([[True, True], [True, False], [False, False]])
print(t.numpy())

[[ True  True]
 [ True False]
 [False False]]

print(t.all().numpy())

False

print(t.all(axis=0).numpy())

[False False]

print(t.all(axis=1, keepdim=True).numpy())

[[ True]
 [False]
 [False]]

Source code in tinygrad/tensor.py

def all(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
  """
  Tests if all element evaluates to `True` along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the reduce axis and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[True, True], [True, False], [False, False]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.all().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.all(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.all(axis=1, keepdim=True).numpy())
  ```
  """
  return self.bool().min(axis, keepdim)

isclose

isclose(
    other: Tensor,
    rtol: float = 1e-05,
    atol: float = 1e-08,
    equal_nan=False,
) -> Tensor

Returns a new tensor with element-wise comparison of closeness to other within a tolerance.

The rtol and atol keyword arguments control the relative and absolute tolerance of the comparison.

By default, two NaN values are not close to each other. If equal_nan is True, two NaN values are considered close.

print(Tensor([1e-7, 1e-8, 1e-9, float('nan')]).isclose(Tensor([0.0, 0.0, 0.0, float('nan')])).numpy())

[False  True  True False]

print(Tensor([float('nan')]).isclose(Tensor([float('nan')]), equal_nan=True).numpy())

[ True]

Source code in tinygrad/tensor.py

def isclose(self, other:Tensor, rtol:float=1e-05, atol:float=1e-08, equal_nan=False) -> Tensor:
  """
  Returns a new tensor with element-wise comparison of closeness to `other` within a tolerance.

  The `rtol` and `atol` keyword arguments control the relative and absolute tolerance of the comparison.

  By default, two `NaN` values are not close to each other. If `equal_nan` is `True`, two `NaN` values are considered close.

  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor([1e-7, 1e-8, 1e-9, float('nan')]).isclose(Tensor([0.0, 0.0, 0.0, float('nan')])).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor([float('nan')]).isclose(Tensor([float('nan')]), equal_nan=True).numpy())
  ```
  """
  is_finite_close = self.isfinite() & other.isfinite() & ((self - other).abs() <= atol + rtol * other.abs())
  is_infinite_close = (self.isinf() | other.isinf()) & (self == other)
  is_nan_close = (self.isnan() & other.isnan()) & equal_nan
  return is_finite_close | is_infinite_close | is_nan_close

allclose

allclose(
    other: Tensor,
    rtol: float = 1e-05,
    atol: float = 1e-08,
    equal_nan=False,
) -> bool

Check if all self and other are close. Return True or False.

Source code in tinygrad/tensor.py

def allclose(self, other:Tensor, rtol:float=1e-05, atol:float=1e-08, equal_nan=False) -> bool:
  """
  Check if all self and other are close. Return True or False.
  """
  return bool(self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan).all().item())

mean

mean(
    axis: int | Sequence[int] | None = None, keepdim=False
) -> Tensor

Returns the mean value of the tensor along the specified axis or axes.

You can pass in axis and keepdim keyword arguments to control the axis along which the mean is computed and whether the reduced dimensions are retained.

Tensor.manual_seed(42)
t = Tensor.normal(2, 3, mean=2.5, std=0.5)
print(t.numpy())

[[3.4788 2.407  3.3202]
 [2.1177 2.0653 2.2811]]

print(t.mean().numpy())

2.6116748

print(t.mean(axis=0).numpy())

[2.7982 2.2361 2.8006]

print(t.mean(axis=1).numpy())

[3.0687 2.1547]

Source code in tinygrad/tensor.py

def mean(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Tensor:
  """
  Returns the mean value of the tensor along the specified axis or axes.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the mean is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.normal(2, 3, mean=2.5, std=0.5)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.mean().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.mean(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.mean(axis=1).numpy())
  ```
  """
  output_dtype = self.dtype if dtypes.is_float(self.dtype) else dtypes.float32
  numerator = self.cast(sum_acc_dtype(self.dtype)).sum(axis=axis, keepdim=keepdim)
  denominator = prod([si for si, so in zip(self.shape, self.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])
  return numerator.div(denominator).cast(output_dtype)

var

var(
    axis: int | Sequence[int] | None = None,
    keepdim=False,
    correction=1,
) -> Tensor

Returns the variance of the tensor along the specified axis or axes.

You can pass in axis, keepdim, and correction keyword arguments to control the axis along which the variance is computed, whether the reduced dimensions are retained, and the Bessel's correction applied.

Tensor.manual_seed(42)
t = Tensor.normal(2, 3, mean=2.5, std=0.5)
print(t.numpy())

[[3.4788 2.407  3.3202]
 [2.1177 2.0653 2.2811]]

print(t.var().numpy())

0.38955206

print(t.var(axis=0).numpy())

[0.9264 0.0584 0.5399]

print(t.var(axis=1).numpy())

[0.3346 0.0127]

Source code in tinygrad/tensor.py

def var(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> Tensor:
  """
  Returns the variance of the tensor along the specified axis or axes.

  You can pass in `axis`, `keepdim`, and `correction` keyword arguments to control the axis along
  which the variance is computed, whether the reduced dimensions are retained, and the Bessel's correction applied.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.normal(2, 3, mean=2.5, std=0.5)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.var().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.var(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.var(axis=1).numpy())
  ```
  """
  squares = (self - self.mean(axis=axis, keepdim=True)).square()
  n = prod([si for si, so in zip(self.shape, squares.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])
  denominator = Tensor(n, device=self.device) - correction
  # TODO: infer device and remove relu
  return squares.sum(axis=axis, keepdim=keepdim).div(denominator.relu())

var_mean

var_mean(
    axis: int | Sequence[int] | None = None,
    keepdim=False,
    correction=1,
) -> tuple[Tensor, Tensor]

Calculates the variance and mean over the dimensions specified by dim. Syntactic sugar around Tensor.var and Tensor.mean to match torch.var_mean.

Tensor.manual_seed(42)
t = Tensor.normal(2, 3, mean=2.5, std=0.5)
print(t.numpy())

[[3.4788 2.407  3.3202]
 [2.1177 2.0653 2.2811]]

var, mean = t.var_mean()
print(var.numpy(), mean.numpy())

0.38955206 2.6116748

Source code in tinygrad/tensor.py

def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Tensor, Tensor]:
  """
  Calculates the variance and mean over the dimensions specified by dim.
  Syntactic sugar around `Tensor.var` and `Tensor.mean` to match `torch.var_mean`.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.normal(2, 3, mean=2.5, std=0.5)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  var, mean = t.var_mean()
  print(var.numpy(), mean.numpy())
  ```
  """
  return self.var(axis, keepdim, correction), self.mean(axis, keepdim)

std

std(
    axis: int | Sequence[int] | None = None,
    keepdim=False,
    correction=1,
) -> Tensor

Returns the standard deviation of the tensor along the specified axis or axes.

You can pass in axis, keepdim, and correction keyword arguments to control the axis along which the standard deviation is computed, whether the reduced dimensions are retained, and the Bessel's correction applied.

Tensor.manual_seed(42)
t = Tensor.normal(2, 3, mean=2.5, std=0.5)
print(t.numpy())

[[3.4788 2.407  3.3202]
 [2.1177 2.0653 2.2811]]

print(t.std().numpy())

0.62414104

print(t.std(axis=0).numpy())

[0.9625 0.2417 0.7348]

print(t.std(axis=1).numpy())

[0.5785 0.1126]

Source code in tinygrad/tensor.py

def std(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> Tensor:
  """
  Returns the standard deviation of the tensor along the specified axis or axes.

  You can pass in `axis`, `keepdim`, and `correction` keyword arguments to control the axis along
  which the standard deviation is computed, whether the reduced dimensions are retained, and the Bessel's correction applied.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.normal(2, 3, mean=2.5, std=0.5)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.std().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.std(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.std(axis=1).numpy())
  ```
  """
  return self.var(axis, keepdim, correction).sqrt()

std_mean

std_mean(
    axis: int | Sequence[int] | None = None,
    keepdim=False,
    correction=1,
) -> tuple[Tensor, Tensor]

Calculates the standard deviation and mean over the dimensions specified by dim. Syntactic sugar around Tensor.std and Tensor.mean to match torch.std_mean.

Tensor.manual_seed(42)
t = Tensor.normal(2, 3, mean=2.5, std=0.5)
print(t.numpy())

[[3.4788 2.407  3.3202]
 [2.1177 2.0653 2.2811]]

std, mean = t.std_mean()
print(std.numpy(), mean.numpy())

0.62414104 2.6116748

Source code in tinygrad/tensor.py

def std_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Tensor, Tensor]:
  """
  Calculates the standard deviation and mean over the dimensions specified by dim.
  Syntactic sugar around `Tensor.std` and `Tensor.mean` to match `torch.std_mean`.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.normal(2, 3, mean=2.5, std=0.5)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  std, mean = t.std_mean()
  print(std.numpy(), mean.numpy())
  ```
  """
  return self.std(axis, keepdim, correction), self.mean(axis, keepdim)

softmax

softmax(axis=-1, dtype: DTypeLike | None = None) -> Tensor

Applies the softmax function to the tensor along the specified axis.

Rescales the elements of the tensor such that they lie in the range [0, 1] and sum to 1.

You can pass in the axis keyword argument to control the axis along which the softmax is computed.

Tensor.manual_seed(42)
t = Tensor.randn(2, 3)
print(t.numpy())

[[ 1.9576 -0.1859  1.6404]
 [-0.7647 -0.8695 -0.4379]]

print(t.softmax().numpy())

[[0.5419 0.0635 0.3946]
 [0.3042 0.274  0.4218]]

print(t.softmax(axis=0).numpy())

[[0.9383 0.6645 0.8888]
 [0.0617 0.3355 0.1112]]

Source code in tinygrad/tensor.py

def softmax(self, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
  """
  Applies the softmax function to the tensor along the specified axis.

  Rescales the elements of the tensor such that they lie in the range [0, 1] and sum to 1.

  You can pass in the `axis` keyword argument to control the axis along which the softmax is computed.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.randn(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.softmax().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.softmax(axis=0).numpy())
  ```
  """
  _, e, ss = self._softmax(axis, dtype)
  return e.div(ss)

log_softmax

log_softmax(
    axis=-1, dtype: DTypeLike | None = None
) -> Tensor

Applies the log-softmax function to the tensor along the specified axis.

The log-softmax function is a numerically stable alternative to the softmax function in log space.

You can pass in the axis keyword argument to control the axis along which the log-softmax is computed.

Tensor.manual_seed(42)
t = Tensor.randn(2, 3)
print(t.numpy())

[[ 1.9576 -0.1859  1.6404]
 [-0.7647 -0.8695 -0.4379]]

print(t.log_softmax().numpy())

[[-0.6127 -2.7563 -0.9299]
 [-1.19   -1.2948 -0.8632]]

print(t.log_softmax(axis=0).numpy())

[[-0.0637 -0.4087 -0.1179]
 [-2.786  -1.0922 -2.1962]]

Source code in tinygrad/tensor.py

def log_softmax(self, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
  """
  Applies the log-softmax function to the tensor along the specified axis.

  The log-softmax function is a numerically stable alternative to the softmax function in log space.

  You can pass in the `axis` keyword argument to control the axis along which the log-softmax is computed.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.randn(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.log_softmax().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.log_softmax(axis=0).numpy())
  ```
  """
  m, _, ss = self._softmax(axis, dtype)
  return m - ss.log()

logsumexp

logsumexp(axis=None, keepdim=False) -> Tensor

Computes the log-sum-exp of the tensor along the specified axis or axes.

The log-sum-exp function is a numerically stable way to compute the logarithm of the sum of exponentials.

You can pass in axis and keepdim keyword arguments to control the axis along which the log-sum-exp is computed and whether the reduced dimensions are retained.

Tensor.manual_seed(42)
t = Tensor.randn(2, 3)
print(t.numpy())

[[ 1.9576 -0.1859  1.6404]
 [-0.7647 -0.8695 -0.4379]]

print(t.logsumexp().numpy())

2.681043

print(t.logsumexp(axis=0).numpy())

[2.0213 0.2227 1.7583]

print(t.logsumexp(axis=1).numpy())

[2.5703 0.4253]

Source code in tinygrad/tensor.py

def logsumexp(self, axis=None, keepdim=False) -> Tensor:
  """
  Computes the log-sum-exp of the tensor along the specified axis or axes.

  The log-sum-exp function is a numerically stable way to compute the logarithm of the sum of exponentials.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the log-sum-exp is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.randn(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logsumexp().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logsumexp(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logsumexp(axis=1).numpy())
  ```
  """
  m = self.max(axis=axis, keepdim=True)
  return (self - m).exp().sum(axis=axis, keepdim=keepdim).log() + (m if keepdim else m.squeeze(axis))

logcumsumexp

logcumsumexp(axis=0) -> Tensor

Computes the log-cumsum-exp of the tensor along the specified axis or axes.

The log-cumsum-exp function is a numerically stable way to compute the logarithm of the cumulative sum of exponentials.

You can pass in the axis keyword argument to control the axis along which the log-cumsum-exp is computed.

Tensor.manual_seed(42)
t = Tensor.randn(2, 3)
print(t.numpy())

[[ 1.9576 -0.1859  1.6404]
 [-0.7647 -0.8695 -0.4379]]

print(t.logcumsumexp().numpy())

[[ 1.9576 -0.1859  1.6404]
 [ 2.0213  0.2227  1.7583]]

print(t.logcumsumexp(axis=0).numpy())

[[ 1.9576 -0.1859  1.6404]
 [ 2.0213  0.2227  1.7583]]

print(t.logcumsumexp(axis=1).numpy())

[[ 1.9576  2.0685  2.5703]
 [-0.7647 -0.1226  0.4253]]

Source code in tinygrad/tensor.py

def logcumsumexp(self, axis=0) -> Tensor:
  """
  Computes the log-cumsum-exp of the tensor along the specified axis or axes.

  The log-cumsum-exp function is a numerically stable way to compute the logarithm of the cumulative sum of exponentials.

  You can pass in the `axis` keyword argument to control the axis along which
  the log-cumsum-exp is computed.

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.randn(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logcumsumexp().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logcumsumexp(axis=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.logcumsumexp(axis=1).numpy())
  ```
  """
  if self.ndim == 0: return self
  x = self.transpose(axis, -1)
  last_dim_size = x.shape[-1]
  x_unsqueezed = x.unsqueeze(-2).expand((None,)*(self.ndim-1)+(last_dim_size, None))
  x_cummax, _ = x.cummax(-1)
  mask = Tensor.ones(last_dim_size, last_dim_size, requires_grad=False, device=self.device).tril()
  ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax
  return ret.transpose(-1, axis)

argmax

argmax(axis=None, keepdim=False) -> Tensor

Returns the indices of the maximum value of the tensor along the specified axis.

You can pass in axis and keepdim keyword arguments to control the axis along which the maximum is computed and whether the reduced dimensions are retained.

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

[[1 0 2]
 [5 4 3]]

print(t.argmax().numpy()) # Returns the index of the maximum value in the flattened tensor.

3

print(t.argmax(axis=0).numpy()) # Returns the indices of the maximum values along axis 0.

[1 1 1]

print(t.argmax(axis=1).numpy()) # Returns the indices of the maximum values along axis 1.

[2 0]

Source code in tinygrad/tensor.py

def argmax(self, axis=None, keepdim=False) -> Tensor:
  """
  Returns the indices of the maximum value of the tensor along the specified axis.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the maximum is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 0, 2], [5, 4, 3]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmax().numpy()) # Returns the index of the maximum value in the flattened tensor.
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmax(axis=0).numpy()) # Returns the indices of the maximum values along axis 0.
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmax(axis=1).numpy()) # Returns the indices of the maximum values along axis 1.
  ```
  """
  if axis is None: return self.flatten().argmax(0)
  axis = self._resolve_dim(axis)
  m = self == self.max(axis=axis, keepdim=True)
  idx = m * Tensor.arange(self.shape[axis],0,-1, requires_grad=False, device=self.device).reshape(self.shape[axis], *[1]*(self.ndim-axis-1))
  return (self.shape[axis]-idx.max(axis=axis, keepdim=keepdim)).cast(dtypes.int32)

argmin

argmin(axis=None, keepdim=False) -> Tensor

Returns the indices of the minimum value of the tensor along the specified axis.

You can pass in axis and keepdim keyword arguments to control the axis along which the minimum is computed and whether the reduced dimensions are retained.

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

[[1 0 2]
 [5 4 3]]

print(t.argmin().numpy()) # Returns the index of the minimum value in the flattened tensor.

1

print(t.argmin(axis=0).numpy()) # Returns the indices of the minimum values along axis 0.

[0 0 0]

print(t.argmin(axis=1).numpy()) # Returns the indices of the minimum values along axis 1.

[1 2]

Source code in tinygrad/tensor.py

def argmin(self, axis=None, keepdim=False) -> Tensor:
  """
  Returns the indices of the minimum value of the tensor along the specified axis.

  You can pass in `axis` and `keepdim` keyword arguments to control the axis along
  which the minimum is computed and whether the reduced dimensions are retained.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 0, 2], [5, 4, 3]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmin().numpy()) # Returns the index of the minimum value in the flattened tensor.
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmin(axis=0).numpy()) # Returns the indices of the minimum values along axis 0.
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.argmin(axis=1).numpy()) # Returns the indices of the minimum values along axis 1.
  ```
  """
  return self._inverse().argmax(axis=axis, keepdim=keepdim)

Processing

avg_pool2d

avg_pool2d(
    kernel_size: tuple[int, ...] = (2, 2),
    stride=None,
    dilation=1,
    padding: int | tuple[int, ...] = 0,
    ceil_mode=False,
    count_include_pad=True,
) -> Tensor

Applies average pooling over a tensor.

This function supports three different types of padding

1. int (single value): Applies the same padding value uniformly to all spatial dimensions.

2. tuple[int, ...] (length = number of spatial dimensions): Specifies a distinct padding value for each spatial dimension in the form (padding_height, padding_width, ...).

3. tuple[int, ...] (length = 2 * number of spatial dimensions): Specifies explicit padding for each side of each spatial dimension in the form (padding_left, padding_right, padding_top, padding_bottom, ...).

When ceil_mode is set to True, output shape will be determined using ceil division. When count_include_pad is set to False, zero padding will not be included in the averaging calculation.

Note

unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.

t = Tensor.arange(25).reshape(1, 1, 5, 5)
print(t.avg_pool2d().numpy())

[[[[ 3.  5.]
   [13. 15.]]]]

print(t.avg_pool2d(ceil_mode=True).numpy())

[[[[ 3.   5.   6.5]
   [13.  15.  16.5]
   [20.5 22.5 24. ]]]]

print(t.avg_pool2d(padding=1).numpy())

[[[[ 0.    0.75  1.75]
   [ 3.75  9.   11.  ]
   [ 8.75 19.   21.  ]]]]

print(t.avg_pool2d(padding=1, count_include_pad=False).numpy())

[[[[ 0.   1.5  3.5]
   [ 7.5  9.  11. ]
   [17.5 19.  21. ]]]]

Source code in tinygrad/tensor.py

def avg_pool2d(self, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0,
               ceil_mode=False, count_include_pad=True) -> Tensor:
  """
  Applies average pooling over a tensor.

  This function supports three different types of `padding`

  1. `int` (single value):
    Applies the same padding value uniformly to all spatial dimensions.

  2. `tuple[int, ...]` (length = number of spatial dimensions):
    Specifies a distinct padding value for each spatial dimension in the form `(padding_height, padding_width, ...)`.

  3. `tuple[int, ...]` (length = 2 * number of spatial dimensions):
    Specifies explicit padding for each side of each spatial dimension in the form
    `(padding_left, padding_right, padding_top, padding_bottom, ...)`.

  When `ceil_mode` is set to `True`, output shape will be determined using ceil division.
  When `count_include_pad` is set to `False`, zero padding will not be included in the averaging calculation.

  NOTE: unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(25).reshape(1, 1, 5, 5)
  print(t.avg_pool2d().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.avg_pool2d(ceil_mode=True).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.avg_pool2d(padding=1).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.avg_pool2d(padding=1, count_include_pad=False).numpy())
  ```
  """
  axis = tuple(range(-len(k_ := make_tuple(kernel_size, 2)), 0))
  def pool(x:Tensor, padding_:Sequence[int]) -> Tensor: return x.pad(padding_)._pool(k_, stride if stride is not None else k_, dilation)
  reg_pads = self._resolve_pool_pads(padding, len(k_))
  ceil_pads = self._apply_ceil_mode(reg_pads, k_, stride if stride is not None else k_, dilation)
  if not count_include_pad:
    pads = ceil_pads if ceil_mode else reg_pads
    return pool(self, pads).sum(axis) / pool(self.ones_like(), pads).sum(axis)
  if not ceil_mode: return pool(self, reg_pads).mean(axis)
  return pool(self, ceil_pads).sum(axis) / pool(self.pad(reg_pads).ones_like(), tuple(cp-rp for cp,rp in zip(ceil_pads, reg_pads))).sum(axis)

max_pool2d

max_pool2d(
    kernel_size: tuple[int, ...] = (2, 2),
    stride=None,
    dilation=1,
    padding: int | tuple[int, ...] = 0,
    ceil_mode=False,
    return_indices=False,
) -> Tensor | tuple[Tensor, Tensor]

Applies max pooling over a tensor.

This function supports three different types of padding

1. int (single value): Applies the same padding value uniformly to all spatial dimensions.

2. tuple[int, ...] (length = number of spatial dimensions): Specifies a distinct padding value for each spatial dimension in the form (padding_height, padding_width, ...).

3. tuple[int, ...] (length = 2 * number of spatial dimensions): Specifies explicit padding for each side of each spatial dimension in the form (padding_left, padding_right, padding_top, padding_bottom, ...).

When ceil_mode is set to True, output shape will be determined using ceil division. When return_indices is set to True, the argmax will be returned along with the max values.

Note

unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.

t = Tensor.arange(25).reshape(1, 1, 5, 5)
print(t.max_pool2d().numpy())

[[[[ 6  8]
   [16 18]]]]

print(t.max_pool2d(ceil_mode=True).numpy())

[[[[ 6  8  9]
   [16 18 19]
   [21 23 24]]]]

print(t.max_pool2d(padding=1).numpy())

[[[[ 0  2  4]
   [10 12 14]
   [20 22 24]]]]

Source code in tinygrad/tensor.py

def max_pool2d(self, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0,
               ceil_mode=False, return_indices=False) -> Tensor | tuple[Tensor, Tensor]:
  """
  Applies max pooling over a tensor.

  This function supports three different types of `padding`

  1. `int` (single value):
    Applies the same padding value uniformly to all spatial dimensions.

  2. `tuple[int, ...]` (length = number of spatial dimensions):
    Specifies a distinct padding value for each spatial dimension in the form `(padding_height, padding_width, ...)`.

  3. `tuple[int, ...]` (length = 2 * number of spatial dimensions):
    Specifies explicit padding for each side of each spatial dimension in the form
    `(padding_left, padding_right, padding_top, padding_bottom, ...)`.

  When `ceil_mode` is set to `True`, output shape will be determined using ceil division.
  When `return_indices` is set to `True`, the argmax will be returned along with the max values.

  NOTE: unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(25).reshape(1, 1, 5, 5)
  print(t.max_pool2d().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.max_pool2d(ceil_mode=True).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.max_pool2d(padding=1).numpy())
  ```
  """
  axis = tuple(range(-len(k_ := make_tuple(kernel_size, 2)), 0))
  pads = self._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)
  pooled = self.pad(pads, value=self.dtype.min)._pool(k_, stride if stride is not None else k_, dilation)
  if not return_indices: return pooled.max(axis)
  spatial_sz = int(math.prod(spatial_shape := self.shape[-len(k_):]))
  idx = Tensor.arange(spatial_sz,0,-1, requires_grad=False, device=self.device).reshape(spatial_shape)
  m = pooled == pooled.max(axis, keepdim=True)
  idx = m * idx.pad(pads, value=idx.dtype.min)._pool(k_, stride if stride is not None else k_, dilation)
  return pooled.max(axis), spatial_sz - idx.max(axis)

max_unpool2d

max_unpool2d(
    indices: Tensor,
    kernel_size: tuple[int, ...] = (2, 2),
    stride=None,
    dilation=1,
    padding: int | tuple[int, ...] = 0,
    output_size=None,
)

Performs a partial inverse of max_pool2d using the indices from the argmax.

When output_size is provided, the output shape disambiguates to the provided shape.

Note

unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.

t = Tensor.arange(1, 17).reshape(1, 1, 4, 4)
print(t.numpy())

[[[[ 1  2  3  4]
   [ 5  6  7  8]
   [ 9 10 11 12]
   [13 14 15 16]]]]

output, indices = Tensor.max_pool2d(t, return_indices=True)
print(output.numpy())
print(indices.numpy())

[[[[ 6  8]
   [14 16]]]]
[[[[ 5  7]
   [13 15]]]]

print(Tensor.max_unpool2d(output, indices).numpy())

[[[[ 0  0  0  0]
   [ 0  6  0  8]
   [ 0  0  0  0]
   [ 0 14  0 16]]]]

Source code in tinygrad/tensor.py

def max_unpool2d(self, indices:Tensor, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0, output_size=None):
  """
  Performs a partial inverse of `max_pool2d` using the indices from the argmax.

  When `output_size` is provided, the output shape disambiguates to the provided shape.

  NOTE: unlike PyTorch, this implementation is not limited to only 2d pooling and instead works for any number of dimensions.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(1, 17).reshape(1, 1, 4, 4)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  output, indices = Tensor.max_pool2d(t, return_indices=True)
  print(output.numpy())
  print(indices.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.max_unpool2d(output, indices).numpy())
  ```
  """
  bs,c,*spatial_shape = self.shape
  if output_size is None:
    k_,d_,s_ = (make_tuple(x, len(spatial_shape)) for x in (kernel_size, dilation, stride if stride is not None else kernel_size))
    p_ = _flat_to_grouped(self._resolve_pool_pads(padding, len(spatial_shape)))
    # https://arxiv.org/pdf/1603.07285 inverse of relationship 15 in section 5.1.
    output_size = tuple((i-1)*s - (pB+pA) + (d*(k-1)+1) for i,k,d,s,(pA,pB) in zip(spatial_shape,k_,d_,s_,p_))
  else: output_size = output_size[-len(spatial_shape):]
  ret = (indices.reshape(bs,c,1,-1)._one_hot_along_dim(prod(output_size), 2).where(self.reshape(bs,c,1,-1), 0)).sum(3)
  return ret.reshape(bs,c,*output_size)

conv2d

conv2d(
    weight: Tensor,
    bias: Tensor | None = None,
    groups=1,
    stride=1,
    dilation=1,
    padding: int | tuple[int, ...] = 0,
    dtype: DTypeLike | None = None,
) -> Tensor

Applies a convolution over a tensor with a given weight and optional bias.

This function supports three different types of padding

1. int (single value): Applies the same padding value uniformly to all spatial dimensions.

2. tuple[int, ...] (length = number of spatial dimensions): Specifies a distinct padding value for each spatial dimension in the form (padding_height, padding_width, ...).

3. tuple[int, ...] (length = 2 * number of spatial dimensions): Specifies explicit padding for each side of each spatial dimension in the form (padding_left, padding_right, padding_top, padding_bottom, ...).

Note

unlike PyTorch, this implementation is not limited to only 2d convolutions and instead works for any number of dimensions.

See: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html

t = Tensor.arange(9).reshape(1, 1, 3, 3)
w = Tensor.ones(1, 1, 2, 2)
print(t.conv2d(w).numpy())

[[[[ 8. 12.]
   [20. 24.]]]]

Source code in tinygrad/tensor.py

def conv2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding:int|tuple[int, ...]=0,
           dtype:DTypeLike|None=None) -> Tensor:
  """
  Applies a convolution over a tensor with a given `weight` and optional `bias`.

  This function supports three different types of `padding`

  1. `int` (single value):
    Applies the same padding value uniformly to all spatial dimensions.

  2. `tuple[int, ...]` (length = number of spatial dimensions):
    Specifies a distinct padding value for each spatial dimension in the form `(padding_height, padding_width, ...)`.

  3. `tuple[int, ...]` (length = 2 * number of spatial dimensions):
    Specifies explicit padding for each side of each spatial dimension in the form
    `(padding_left, padding_right, padding_top, padding_bottom, ...)`.

  NOTE: unlike PyTorch, this implementation is not limited to only 2d convolutions and instead works for any number of dimensions.

  See: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(9).reshape(1, 1, 3, 3)
  w = Tensor.ones(1, 1, 2, 2)
  print(t.conv2d(w).numpy())
  ```
  """
  if IMAGE: return self.image_conv2d(weight, bias, groups, stride, dilation, padding, dtype)
  (bs,cin_), (cout,cin), HW = self.shape[:2], weight.shape[:2], weight.shape[2:]
  padding_ = self._resolve_pool_pads(padding, len(HW))
  assert groups*cin == cin_ and len(self.shape) == len(weight.shape),\
      f"Input Tensor shape {self.shape} does not match the shape of the weights {weight.shape}. ({groups*cin} vs. {cin_})"

  # conv2d is a pooling op (with padding)
  x = self.pad(padding_)._pool(HW, stride, dilation)   # (bs, groups*cin, oy, ox, H, W)
  rcout, oyx = cout//groups, x.shape[2:-len(HW)]
  if not all(x == 3 for x in HW) or stride != 1 or dilation != 1 or not WINO:
    # normal conv
    x = x.reshape(bs, groups, cin, 1, *oyx, *HW).expand(bs, groups, cin, rcout, *oyx, *HW)\
      .permute(0,1,3,*[4+i for i in range(len(oyx))],2,*[4+len(oyx)+i for i in range(len(HW))])

    # conv! broadcasted to (bs, groups, rcout, *oyx, cin, *HW)
    ret = (x * weight.reshape(1, groups, rcout, *[1] * len(oyx), cin, *HW))\
      .sum([-1-i for i in range(1+len(oyx))], keepdim=True, dtype=dtype).reshape(bs, cout, *oyx)
    return ret if bias is None else ret.add(bias.reshape(1, -1, *[1] * len(HW)))

  HWI, HWO = (6,) * len(HW), (4,) * len(HW)  # F(4x4,3x3) winograd tiles
  winograd_G = [[1/4, 0, 0], [-1/6, -1/6, -1/6], [-1/6, 1/6, -1/6], [1/24, 1/12, 1/6], [1/24, -1/12, 1/6], [0, 0, 1]]
  winograd_Bt = [[4, 0, -5, 0, 1, 0], [0, -4, -4, 1, 1, 0], [0, 4, -4, -1, 1, 0], [0, -2, -1, 2, 1, 0], [0, 2, -1, -2, 1, 0], [0, 4, 0, -5, 0, 1]]
  winograd_At = [[1, 1, 1, 1, 1, 0], [0, 1, -1, 2, -2, 0], [0, 1, 1, 4, 4, 0], [0, 1, -1, 8, -8, 1]] # applying At in pre-order doubles compile time

  # TODO: stride == dilation
  # use padding to round up to 4x4 output tiles
  # (bs, cin_, tyx, HWI)
  pads = [[padding_[i*2], padding_[i*2+1] + (-(dim+sum(padding_[i*2:(i+1)*2])-2) % 4)] for i, dim in enumerate(reversed(self.shape[-len(HW):]))]
  d = self.pad(sum(pads, []))._pool(HWI, HWO)
  # move HW to the front: # (HWI, bs, cin_, tyx)
  d = d.permute(*range(len(d.shape)-len(HW),len(d.shape)), *range(len(d.shape)-len(HW)))
  tyx = d.shape[-len(HWI):]  # dim of tiling

  g = weight.permute(*range(len(weight.shape)-len(HW),len(weight.shape)), *range(len(weight.shape)-len(HW)))  # move HW to the front

  # compute 6x6 winograd tiles: GgGt, BtdB
  # (HWI, groups * rcout, cin) -> (HWI, bs=1, groups, rcout, cin, tyx=(1,1))
  gfactors = _apply_winograd_matrix(winograd_G, g, len(HW)).reshape(*HWI, 1, groups, rcout, cin, *([1]*len(tyx)))
  # (HWI, bs, cin_, tyx) -> (HWI, bs, groups, 1 ,cin, *tyx)
  dfactors = _apply_winograd_matrix(winograd_Bt, d, len(HW)).reshape(*HWI, bs, groups, 1, cin, *tyx)

  # matmul; sum across cin: (HWI, bs, groups, rcout, *tyx); then HWI -> HWO: (HWO, bs, groups, rcout, *tyx)
  ret = _apply_winograd_matrix(winograd_At, (gfactors * dfactors).sum(axis=-1-len(HW), dtype=dtype), len(HW))

  # interleave tyx and HWO: (bs, groups, rcout, oy, HO, ox, WO)
  ret = ret.permute([*range(len(HW), len(ret.shape)-len(HW)), *[i+o for i in range(len(HW)) for o in [len(ret.shape)-len(HW),0]]])
  # merge groups and rcout, tyx and HWO: (bs, groups, cout, *yx), shrink to final
  ret = ret.reshape(bs, cout, *[c * HWO[i] for i, c in enumerate(tyx)]).shrink_to(bs, cout, *oyx)

  return (ret if bias is None else ret.add(bias.reshape(1, -1, *[1 for _ in range(len(HW))]))).contiguous().contiguous_backward()

conv_transpose2d

conv_transpose2d(
    weight: Tensor,
    bias: Tensor | None = None,
    groups=1,
    stride=1,
    dilation=1,
    padding=0,
    output_padding=0,
) -> Tensor

Applies a transposed convolution over a tensor with a given weight and optional bias.

This function supports three different types of padding

1. int (single value): Applies the same padding value uniformly to all spatial dimensions.

2. tuple[int, ...] (length = number of spatial dimensions): Specifies a distinct padding value for each spatial dimension in the form (padding_height, padding_width, ...).

3. tuple[int, ...] (length = 2 * number of spatial dimensions): Specifies explicit padding for each side of each spatial dimension in the form (padding_left, padding_right, padding_top, padding_bottom, ...).

Note

unlike PyTorch, this implementation is not limited to only 2d transposed convolutions and instead works for any number of dimensions.

See: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html

t = Tensor.arange(9).reshape(1, 1, 3, 3)
w = Tensor.ones(1, 1, 2, 2)
print(t.conv_transpose2d(w).numpy())

[[[[ 0.  1.  3.  2.]
   [ 3.  8. 12.  7.]
   [ 9. 20. 24. 13.]
   [ 6. 13. 15.  8.]]]]

Source code in tinygrad/tensor.py

def conv_transpose2d(self, weight:Tensor, bias:Tensor|None=None, groups=1, stride=1, dilation=1, padding=0, output_padding=0) -> Tensor:
  """
  Applies a transposed convolution over a tensor with a given `weight` and optional `bias`.

  This function supports three different types of `padding`

  1. `int` (single value):
    Applies the same padding value uniformly to all spatial dimensions.

  2. `tuple[int, ...]` (length = number of spatial dimensions):
    Specifies a distinct padding value for each spatial dimension in the form `(padding_height, padding_width, ...)`.

  3. `tuple[int, ...]` (length = 2 * number of spatial dimensions):
    Specifies explicit padding for each side of each spatial dimension in the form
    `(padding_left, padding_right, padding_top, padding_bottom, ...)`.

  NOTE: unlike PyTorch, this implementation is not limited to only 2d transposed convolutions and instead works for any number of dimensions.

  See: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(9).reshape(1, 1, 3, 3)
  w = Tensor.ones(1, 1, 2, 2)
  print(t.conv_transpose2d(w).numpy())
  ```
  """
  x, w = self, weight.unflatten(0, (groups, -1)).transpose(1, 2).flip(*range(3, len(weight.shape)+1))
  HW = weight.shape[2:]
  padding = _flat_to_grouped(self._resolve_pool_pads(padding, len(HW)))
  stride, dilation, output_padding = [make_tuple(x, len(HW)) for x in (stride, dilation, output_padding)]
  if any(s>1 for s in stride):
    # handle strides: (k) -> reshape -> (k,1) -> pad -> (k,s) -> reshape -> (k*s) -> shrink (k-(s-1))
    x = x.reshape(None, None, *flatten((k,1) for k in x.shape[2:]))
    x = x.pad((None, None, *flatten((None,(0,s-1)) for s in stride)))
    x = x.reshape(None, None, *[k*s for k,s in zip(x.shape[2::2], stride)])
    x = x.shrink_to(None, None, *[k-(s-1) for k,s in zip(x.shape[2:], stride)])
  padding = flatten((((k-1)*d-pB,(k-1)*d-pA+op) for k,d,(pB,pA),op in reversed(list(zip(HW, dilation, padding, output_padding)))))
  return x.conv2d(w.flatten(end_dim=1), groups=groups, bias=bias, dilation=dilation, padding=padding)

dot

dot(w: Tensor, dtype: DTypeLike | None = None) -> Tensor

Source code in tinygrad/tensor.py

def dot(self, w:Tensor, dtype:DTypeLike|None=None) -> Tensor:
  if IMAGE: return self.image_dot(w, dtype)
  if ASM_GEMM:
    from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
    if can_use_asm_gemm(self, w): return asm_gemm(self, w)
  return super().dot(w, dtype=dtype)

matmul

matmul(
    x: Self, reverse=False, dtype: DTypeLike | None = None
) -> Self

Performs matrix multiplication between two tensors.

You can pass in the reverse keyword argument to control the order of the matrix multiplication. You can pass in the optional dtype keyword argument to control the data type of the accumulation.

a = Tensor([[1, 2], [3, 4]])
b = Tensor([[5, 6], [7, 8]])
print(a.matmul(b).numpy())

[[19 22]
 [43 50]]

Source code in tinygrad/mixin/__init__.py

def matmul(self, x:Self, reverse=False, dtype:DTypeLike|None=None) -> Self:
  """
  Performs matrix multiplication between two tensors.

  You can pass in the `reverse` keyword argument to control the order of the matrix multiplication.
  You can pass in the optional `dtype` keyword argument to control the data type of the accumulation.

  ```python exec="true" source="above" session="tensor" result="python"
  a = Tensor([[1, 2], [3, 4]])
  b = Tensor([[5, 6], [7, 8]])
  print(a.matmul(b).numpy())
  ```
  """
  return x.dot(self, dtype=dtype) if reverse else self.dot(x, dtype=dtype)

einsum staticmethod

einsum(
    formula: str,
    *operands: Tensor | Sequence[Tensor],
    dtype: DTypeLike | None = None
) -> Tensor

Sums the product of the elements of the input tensors according to a formula based on the Einstein summation convention.

See: https://pytorch.org/docs/stable/generated/torch.einsum.html

x = Tensor([[1, 2], [3, 4]])
y = Tensor([[5, 6], [7, 8]])
print(Tensor.einsum("ij,ij->", x, y).numpy())

70

Source code in tinygrad/tensor.py

@staticmethod
def einsum(formula:str, *operands:Tensor|Sequence[Tensor], dtype:DTypeLike|None=None) -> Tensor:
  """
  Sums the product of the elements of the input tensors according to a formula based on the Einstein summation convention.

  See: https://pytorch.org/docs/stable/generated/torch.einsum.html

  ```python exec="true" source="above" session="tensor" result="python"
  x = Tensor([[1, 2], [3, 4]])
  y = Tensor([[5, 6], [7, 8]])
  print(Tensor.einsum("ij,ij->", x, y).numpy())
  ```
  """
  xs, formula = list(argfix(*operands)), formula.replace(" ", "")
  # expand ellipsis to letters, determine output
  if "..." in formula:
    ell, lhs = "".join(c for c in string.ascii_letters if c not in formula), (formula.split("->") + [""])[0]
    ell_n = [max(0, x.ndim - len(s) + 3) if "..." in s else 0 for s, x in zip(lhs.split(","), xs)]
    for i, (s, x) in enumerate(zip(inputs := lhs.split(","), xs)): inputs[i] = s.replace("...", ell[max(ell_n)-ell_n[i]:max(ell_n)])
    lhs, auto = ",".join(inputs), "".join(sorted(c for c in lhs if lhs.count(c) == 1 and c.isalpha() and c not in ell))
    formula = f"{lhs}->{formula.split('->')[1].replace('...', ell[:max(ell_n)]) if '->' in formula else ell[:max(ell_n)] + auto}"
  lhs, rhs = formula.split("->") if "->" in formula else (formula, "".join(sorted(c for c in formula if formula.count(c)==1 and c.isalpha())))
  inputs = lhs.split(",")
  if len(xs) != len(inputs): raise ValueError(f"number of operands doesn't match, expected {len(inputs)}, got {len(xs)}")
  # trace: take diagonal when letter repeats in single input
  for i, (s, x) in enumerate(zip(inputs, xs)):
    for c in set(s):
      while s.count(c) > 1:
        j, k, n = s.index(c), s.index(c, s.index(c)+1), cast(int, x.shape[s.index(c)])
        perm = [d for d in range(x.ndim) if d not in (j,k)]+[j,k]
        x = x.permute(perm).flatten(-2).pad(((0,0),)*(x.ndim-2)+((0,n),)).unflatten(-1,(n,n+1))[...,0] if x.ndim > 2 else x.diagonal()
        s = s[:k] + s[k+1:]
    inputs[i], xs[i] = s, x
  # check sizes and build sorted alphabet
  sz = merge_dicts([dict(zip(s, x.shape)) for s, x in zip(inputs, xs)])
  alpha = sorted(sz)
  # align all tensors to alphabet, multiply, sum non-output, permute to output order
  xs = [x.permute(*[s.index(c) for c in sorted(s)]).reshape([sz[c] if c in s else 1 for c in alpha]).expand([sz[c] for c in alpha]) if s else x
        for s, x in zip(inputs, xs)]
  return functools.reduce(lambda a,b:a*b, xs).sum([i for i,c in enumerate(alpha) if c not in rhs], dtype=dtype).permute(argsort(argsort(list(rhs))))

cumsum

cumsum(axis: int = 0) -> Tensor

Computes the cumulative sum of the tensor along the specified axis.

t = Tensor.ones(2, 3)
print(t.numpy())

[[1. 1. 1.]
 [1. 1. 1.]]

print(t.cumsum(1).numpy())

[[1. 2. 3.]
 [1. 2. 3.]]

Source code in tinygrad/tensor.py

def cumsum(self, axis:int=0) -> Tensor:
  """
  Computes the cumulative sum of the tensor along the specified `axis`.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.ones(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.cumsum(1).numpy())
  ```
  """
  return self._split_cumalu(axis, Ops.ADD)

cumprod

cumprod(axis: int) -> Tensor

Computes the cumulative product of the elements of the tensor along the specified axis.

t = Tensor.arange(1, 7).reshape(2, 3)
print(t.numpy())

[[1 2 3]
 [4 5 6]]

print(t.cumprod(axis=0).numpy())

[[ 1  2  3]
 [ 4 10 18]]

Source code in tinygrad/tensor.py

def cumprod(self, axis:int) -> Tensor:
  """
  Computes the cumulative product of the elements of the tensor along the specified `axis`.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.arange(1, 7).reshape(2, 3)
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.cumprod(axis=0).numpy())
  ```
  """
  return self._split_cumalu(axis, Ops.MUL)

cummax

cummax(axis: int = 0) -> tuple[Tensor, Tensor]

Computes the cumulative max of the tensor along axis, returning (values, indices).

t = Tensor([0, 1, -1, 2, -2, 3, -3])
values, indices = t.cummax(0)
print(values.numpy())
print(indices.numpy())

[0 1 1 2 2 3 3]
[0 1 1 3 3 5 5]

Source code in tinygrad/tensor.py

def cummax(self, axis:int=0) -> tuple[Tensor, Tensor]:
  """
  Computes the cumulative max of the tensor along `axis`, returning (values, indices).

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([0, 1, -1, 2, -2, 3, -3])
  values, indices = t.cummax(0)
  print(values.numpy())
  print(indices.numpy())
  ```
  """
  if self.ndim == 0: return self._split_cumalu(axis, Ops.MAX), Tensor.zeros(self.shape, dtype=dtypes.int32, device=self.device)
  values, n = self._split_cumalu(axis, Ops.MAX), int(self.shape[axis])
  x, values_t = self.transpose(axis, -1), values.transpose(axis, -1)
  match = (x.unsqueeze(-1) == values_t.unsqueeze(-2)) * Tensor.ones(n, n, requires_grad=False, device=self.device).triu()
  idx = (-(match * Tensor.arange(n, 0, -1, requires_grad=False, device=self.device).reshape(n, 1)).max(-2) + n).cast(dtypes.int32)
  return values, idx.transpose(-1, axis)

cummin

cummin(axis: int = 0) -> tuple[Tensor, Tensor]

Computes the cumulative min of the tensor along axis, returning (values, indices).

t = Tensor([0, 1, -1, 2, -2, 3, -3])
values, indices = t.cummin(0)
print(values.numpy())
print(indices.numpy())

[ 0  0 -1 -1 -2 -2 -3]
[0 0 2 2 4 4 6]

Source code in tinygrad/tensor.py

def cummin(self, axis:int=0) -> tuple[Tensor, Tensor]:
  """
  Computes the cumulative min of the tensor along `axis`, returning (values, indices).

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([0, 1, -1, 2, -2, 3, -3])
  values, indices = t.cummin(0)
  print(values.numpy())
  print(indices.numpy())
  ```
  """
  values, indices = self._inverse().cummax(axis)
  return values._inverse(), indices

triu

triu(diagonal: sint = 0) -> Tensor

Returns the upper triangular part of the tensor, the other elements are set to 0.

The argument diagonal determines which diagonal is on the boundary. diagonal = 0 means the main diagonal. Positive diagonal means above the main diagonal, and negative diagonal means below the main diagonal.

t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
print(t.numpy())

[[ 1  2  3  4]
 [ 5  6  7  8]
 [ 9 10 11 12]]

print(t.triu(diagonal=0).numpy())

[[ 1  2  3  4]
 [ 0  6  7  8]
 [ 0  0 11 12]]

print(t.triu(diagonal=1).numpy())

[[ 0  2  3  4]
 [ 0  0  7  8]
 [ 0  0  0 12]]

print(t.triu(diagonal=-1).numpy())

[[ 1  2  3  4]
 [ 5  6  7  8]
 [ 0 10 11 12]]

Source code in tinygrad/tensor.py

def triu(self, diagonal:sint=0) -> Tensor:
  """
  Returns the upper triangular part of the tensor, the other elements are set to 0.

  The argument `diagonal` determines which diagonal is on the boundary. `diagonal = 0` means the main diagonal.
  Positive `diagonal` means above the main diagonal, and negative `diagonal` means below the main diagonal.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.triu(diagonal=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.triu(diagonal=1).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.triu(diagonal=-1).numpy())
  ```
  """
  return Tensor._tri(self.shape[-2], self.shape[-1], diagonal=diagonal, device=self.device).where(self, self.zeros_like())

tril

tril(diagonal: sint = 0) -> Tensor

Returns the lower triangular part of the tensor, the other elements are set to 0.

The argument diagonal determines which diagonal is on the boundary. diagonal = 0 means the main diagonal. Positive diagonal means above the main diagonal, and negative diagonal means below the main diagonal.

t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
print(t.numpy())

[[ 1  2  3  4]
 [ 5  6  7  8]
 [ 9 10 11 12]]

print(t.tril(diagonal=0).numpy())

[[ 1  0  0  0]
 [ 5  6  0  0]
 [ 9 10 11  0]]

print(t.tril(diagonal=1).numpy())

[[ 1  2  0  0]
 [ 5  6  7  0]
 [ 9 10 11 12]]

print(t.tril(diagonal=-1).numpy())

[[ 0  0  0  0]
 [ 5  0  0  0]
 [ 9 10  0  0]]

Source code in tinygrad/tensor.py

def tril(self, diagonal:sint=0) -> Tensor:
  """
  Returns the lower triangular part of the tensor, the other elements are set to 0.

  The argument `diagonal` determines which diagonal is on the boundary. `diagonal = 0` means the main diagonal.
  Positive `diagonal` means above the main diagonal, and negative `diagonal` means below the main diagonal.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.tril(diagonal=0).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.tril(diagonal=1).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.tril(diagonal=-1).numpy())
  ```
  """
  return Tensor._tri(self.shape[-2], self.shape[-1], diagonal=diagonal+1, device=self.device).where(self.zeros_like(), self)

interpolate

interpolate(
    size: tuple[int, ...],
    mode: str = "linear",
    align_corners: bool = False,
) -> Tensor

Downsamples or Upsamples to the input size, accepts 0 to N batch dimensions.

The interpolation algorithm is selected with mode which currently only supports linear, nearest and nearest-exact. To run bilinear or trilinear, pass in a 2D or 3D size.

t = Tensor([[1, 2, 3, 4], [21, 22, 23, 24], [41, 42, 43, 44]])
print(t.numpy())

[[ 1  2  3  4]
 [21 22 23 24]
 [41 42 43 44]]

print(t.interpolate(size=(2,3), mode="linear").numpy())

[[ 6  7  8]
 [36 37 38]]

Source code in tinygrad/tensor.py

def interpolate(self, size:tuple[int, ...], mode:str="linear", align_corners:bool=False) -> Tensor:
  """
  Downsamples or Upsamples to the input `size`, accepts 0 to N batch dimensions.

  The interpolation algorithm is selected with `mode` which currently only supports `linear`, `nearest` and `nearest-exact`.
  To run `bilinear` or `trilinear`, pass in a 2D or 3D size.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 2, 3, 4], [21, 22, 23, 24], [41, 42, 43, 44]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.interpolate(size=(2,3), mode="linear").numpy())
  ```
  """
  assert isinstance(size, (tuple,list)) and all_int(size) and 0 < len(size) <= self.ndim, f"invalid {size=}"
  assert mode in ("linear", "nearest", "nearest-exact"), "only supports linear, nearest or nearest-exact interpolate"
  assert not (align_corners and mode != "linear"), "align_corners option can only be set with the interpolating mode linear"
  x, expand = self, list(self.shape)
  for i in range(-1,-len(size)-1,-1):
    scale = (int(self.shape[i]) - int(align_corners)) / (size[i] - int(align_corners))
    arr, reshape = Tensor.arange(size[i], dtype=dtypes.float32, device=self.device), [1] * self.ndim
    reshape[i] = expand[i] = size[i]
    if mode == "linear":
      index = (scale*arr if align_corners else (scale*(arr+0.5))-0.5).clip(0, self.shape[i]-1)
      low, high, perc = [y.reshape(reshape).expand(expand) for y in (index.floor().int(), index.ceil().int(), index - index.floor())]
      x = x.gather(i, low).lerp(x.gather(i, high), perc)
    else:
      index = (scale*(arr+0.5) if mode=="nearest-exact" else scale*arr).cast(dtypes.int32).reshape(reshape).expand(expand)
      x = x.gather(i, index)
  return x.cast(self.dtype)

scatter

scatter(
    dim: int,
    index: Tensor,
    src: Tensor | PyConst,
    reduce: Literal["multiply", "add"] | None = None,
) -> Tensor

Scatters src values along an axis specified by dim. Apply add or multiply reduction operation with reduce.

Note

To use the reduce argument with a Tensor src, see Tensor.scatter_reduce.

src = Tensor.arange(1, 11).reshape(2, 5)
print(src.numpy())

[[ 1  2  3  4  5]
 [ 6  7  8  9 10]]

index = Tensor([[0, 1, 2, 0]])
print(Tensor.zeros(3, 5, dtype=src.dtype).scatter(0, index, src).numpy())

[[1 0 0 4 0]
 [0 2 0 0 0]
 [0 0 3 0 0]]

index = Tensor([[0, 1, 2], [0, 1, 4]])
print(Tensor.zeros(3, 5, dtype=src.dtype).scatter(1, index, src).numpy())

[[1 2 3 0 0]
 [6 7 0 0 8]
 [0 0 0 0 0]]

print(Tensor.full((2, 4), 2.0).scatter(1, Tensor([[2], [3]]), 1.23, reduce='multiply').numpy())

[[2.   2.   2.46 2.  ]
 [2.   2.   2.   2.46]]

print(Tensor.full((2, 4), 2.0).scatter(1, Tensor([[2], [3]]), 1.23, reduce='add').numpy())

[[2.   2.   3.23 2.  ]
 [2.   2.   2.   3.23]]

Source code in tinygrad/tensor.py

def scatter(self, dim:int, index:Tensor, src:Tensor|PyConst, reduce:Literal['multiply', 'add']|None=None) -> Tensor:
  """
  Scatters `src` values along an axis specified by `dim`.
  Apply `add` or `multiply` reduction operation with `reduce`.

  NOTE: To use the `reduce` argument with a Tensor `src`, see `Tensor.scatter_reduce`.

  ```python exec="true" source="above" session="tensor" result="python"
  src = Tensor.arange(1, 11).reshape(2, 5)
  print(src.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  index = Tensor([[0, 1, 2, 0]])
  print(Tensor.zeros(3, 5, dtype=src.dtype).scatter(0, index, src).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  index = Tensor([[0, 1, 2], [0, 1, 4]])
  print(Tensor.zeros(3, 5, dtype=src.dtype).scatter(1, index, src).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.full((2, 4), 2.0).scatter(1, Tensor([[2], [3]]), 1.23, reduce='multiply').numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.full((2, 4), 2.0).scatter(1, Tensor([[2], [3]]), 1.23, reduce='add').numpy())
  ```
  """
  if reduce not in {None, "add", "multiply"}: raise TypeError(f"{reduce=} must be one of None, 'multiply', or 'add'")
  if reduce and isinstance(src, Tensor): raise TypeError("Tensor src is not supported with reduce arg. see scatter_reduce")
  if not isinstance(src, Tensor): src = index.full_like(src, device=self.device, dtype=self.dtype)
  if reduce == "add": return self.scatter_reduce(dim, index, src, "sum", include_self=True)
  if reduce == "multiply": return self.scatter_reduce(dim, index, src, "prod", include_self=True)
  src, mask = self._pre_scatter(dim, index, src)
  return _masked_setitem(self, src, mask, (-1,))

scatter_reduce

scatter_reduce(
    dim: int,
    index: Tensor,
    src: Tensor,
    reduce: Literal["sum", "prod", "mean", "amax", "amin"],
    include_self: bool = True,
) -> Tensor

Scatters src values along an axis specified by dim. Apply "sum", "prod", "mean", "amax", or "amin" reduction operations with reduce.

Set include_self=False to exclude values in the self Tensor from the reduction.

src = Tensor.arange(1, 11).cast(dtypes.float).reshape(2, 5)
print(src.numpy())
index = Tensor([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]])
print(index.numpy())

[[ 1.  2.  3.  4.  5.]
 [ 6.  7.  8.  9. 10.]]
[[0 0 0 0 0]
 [0 0 0 0 0]]

print(Tensor.ones(1, 5, dtype=src.dtype).scatter_reduce(0, index, src, reduce='sum').numpy())

[[ 8. 10. 12. 14. 16.]]

print(Tensor.ones(1, 5, dtype=src.dtype).scatter_reduce(0, index, src, reduce='prod').numpy())

[[ 6. 14. 24. 36. 50.]]

print(Tensor.ones(1, 5, dtype=src.dtype).scatter_reduce(0, index, src, reduce='mean', include_self=False).numpy())

[[3.5 4.5 5.5 6.5 7.5]]

print(Tensor([[-10, 20, 0, 5, 10]], dtype=src.dtype).scatter_reduce(0, index, src, reduce='amax').numpy())

[[ 6. 20.  8.  9. 10.]]

print(Tensor([[-10, 20, 0, 5, 10]], dtype=src.dtype).scatter_reduce(0, index, src, reduce='amin').numpy())

[[-10.   2.   0.   4.   5.]]

Source code in tinygrad/tensor.py

def scatter_reduce(self, dim:int, index:Tensor, src:Tensor, reduce:Literal["sum", "prod", "mean", "amax", "amin"],
                   include_self:bool=True) -> Tensor:
  """
  Scatters `src` values along an axis specified by `dim`.
  Apply `"sum"`, `"prod"`, `"mean"`, `"amax"`, or `"amin"` reduction operations with `reduce`.

  Set `include_self=False` to exclude values in the `self` Tensor from the reduction.

  ```python exec="true" source="above" session="tensor" result="python"
  src = Tensor.arange(1, 11).cast(dtypes.float).reshape(2, 5)
  print(src.numpy())
  index = Tensor([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]])
  print(index.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.ones(1, 5, dtype=src.dtype).scatter_reduce(0, index, src, reduce='sum').numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.ones(1, 5, dtype=src.dtype).scatter_reduce(0, index, src, reduce='prod').numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor.ones(1, 5, dtype=src.dtype).scatter_reduce(0, index, src, reduce='mean', include_self=False).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor([[-10, 20, 0, 5, 10]], dtype=src.dtype).scatter_reduce(0, index, src, reduce='amax').numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(Tensor([[-10, 20, 0, 5, 10]], dtype=src.dtype).scatter_reduce(0, index, src, reduce='amin').numpy())
  ```
  """
  src, mask = self._pre_scatter(dim, index, src)
  def _inv_mask(a:Tensor|PyConst, b:Tensor|PyConst) -> Tensor: return mask.any(-1).logical_not().where(a, b)
  if reduce == "sum": return mask.where(src, 0).sum(-1).add(self if include_self else _inv_mask(self, 0))
  if reduce == "prod": return mask.where(src, 1).prod(-1).mul(self if include_self else _inv_mask(self, 1))
  if reduce == "amax": return mask.where(src, m := src.dtype.min).max(-1).maximum(self if include_self else _inv_mask(self, m))
  if reduce == "amin": return mask.where(src, m := src.dtype.max).min(-1).minimum(self if include_self else _inv_mask(self, m))
  if reduce == "mean":
    count = mask.where(1, 0).sum(-1).add(1 if include_self else _inv_mask(1, 0))
    return mask.where(src, 0).sum(-1).add(self if include_self else _inv_mask(self, 0)).div(count)
  raise RuntimeError(f"{reduce=} must be one of 'sum', 'prod', 'mean', 'amax', 'amin'")

masked_select

masked_select(mask)

Selects elements from self based on the boolean mask.

t = Tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
mask = Tensor([[True, False, True], [False, True, False], [False, False, True]])
print(t.numpy())
print(mask.numpy())

[[0 1 2]
 [3 4 5]
 [6 7 8]]
[[ True False  True]
 [False  True False]
 [False False  True]]

print(t.masked_select(mask).numpy())

[0 2 4 8]

Source code in tinygrad/tensor.py

def masked_select(self, mask):
  """
  Selects elements from `self` based on the boolean `mask`.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]])
  mask = Tensor([[True, False, True], [False, True, False], [False, False, True]])
  print(t.numpy())
  print(mask.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.masked_select(mask).numpy())
  ```
  """
  if not dtypes.is_bool(mask.dtype): raise RuntimeError(f"masked_select expects bool mask tensor, got {mask.dtype}")
  x, mask = self.flatten(), mask._broadcast_to(self.shape).flatten()
  mask_cumsum = mask.cumsum()
  counts = Tensor.zeros(mask_cumsum[-1].item(), dtype=dtypes.int32)
  idxs = counts.scatter(0, mask_cumsum, 1, reduce='add').cumsum()
  return x[idxs]

masked_fill

masked_fill(
    mask: Tensor, value: Tensor | PyConst
) -> Tensor

Replaces self with value wherever the elements of mask are True.

t = Tensor([1, 2, 3, 4, 5])
mask = Tensor([True, False, True, False, False])
print(t.masked_fill(mask, -12).numpy())

[-12   2 -12   4   5]

t = Tensor([1, 2, 3, 4, 5])
mask = Tensor([True, False, True, False, False])
value = Tensor([-1, -2, -3, -4, -5])
print(t.masked_fill(mask, value).numpy())

[-1  2 -3  4  5]

Source code in tinygrad/tensor.py

def masked_fill(self:Tensor, mask:Tensor, value:Tensor|PyConst) -> Tensor:
  """
  Replaces `self` with `value` wherever the elements of `mask` are True.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 2, 3, 4, 5])
  mask = Tensor([True, False, True, False, False])
  print(t.masked_fill(mask, -12).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 2, 3, 4, 5])
  mask = Tensor([True, False, True, False, False])
  value = Tensor([-1, -2, -3, -4, -5])
  print(t.masked_fill(mask, value).numpy())
  ```
  """
  return mask.where(value, self)

nonzero

nonzero() -> Tensor

Returns the indices of the elements that are non-zero.

Returns a 2D tensor where each row is the index of a non-zero element.

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

[1 0 2 0 3]

print(t.nonzero().numpy())

[[0]
 [2]
 [4]]

t = Tensor([[1, 0], [0, 2]])
print(t.numpy())

[[1 0]
 [0 2]]

print(t.nonzero().numpy())

[[0 0]
 [1 1]]

Source code in tinygrad/tensor.py

def nonzero(self) -> Tensor:
  """
  Returns the indices of the elements that are non-zero.

  Returns a 2D tensor where each row is the index of a non-zero element.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 0, 2, 0, 3])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.nonzero().numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 0], [0, 2]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.nonzero().numpy())
  ```
  """
  mask = (self != 0).flatten()
  indices = Tensor.stack(*[Tensor.arange(s, device=self.device).reshape(*[1]*i, s, *[1]*(self.ndim-i-1)).expand(self.shape).flatten()
                           for i, s in enumerate(self.shape)], dim=-1)
  return indices.masked_select(mask.unsqueeze(-1).expand(*mask.shape, self.ndim)).reshape(-1, self.ndim)

sort

sort(
    dim: int = -1, descending: bool = False
) -> tuple[Tensor, Tensor]

Performs a bitonic sort on the tensor along the specified dimension.

Order of indices for equivalent elements is always preserved.

See: https://en.wikipedia.org/wiki/Bitonic_sorter

t = Tensor([[0.1, 0.5, 1.2, 3.4, 2.1], [2.2, 1.9, 0.3, 4.5, 0.8]])
print(t.numpy())

[[0.1 0.5 1.2 3.4 2.1]
 [2.2 1.9 0.3 4.5 0.8]]

sorted_values, indices = t.sort(dim=1, descending=True)
print(sorted_values.numpy())
print(indices.numpy())

[[3.4 2.1 1.2 0.5 0.1]
 [4.5 2.2 1.9 0.8 0.3]]
[[3 4 2 1 0]
 [3 0 1 4 2]]

Source code in tinygrad/tensor.py

def sort(self, dim:int=-1, descending:bool=False) -> tuple[Tensor, Tensor]:
  """
  Performs a bitonic sort on the tensor along the specified dimension.

  Order of indices for equivalent elements is always preserved.

  See: https://en.wikipedia.org/wiki/Bitonic_sorter

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[0.1, 0.5, 1.2, 3.4, 2.1], [2.2, 1.9, 0.3, 4.5, 0.8]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  sorted_values, indices = t.sort(dim=1, descending=True)
  print(sorted_values.numpy())
  print(indices.numpy())
  ```
  """
  x, dim = self, self._resolve_dim(dim)
  if (orig_len := int(x.shape[dim])) <= 1: return x, x.zeros_like(dtype=dtypes.default_int)
  # pad to power of 2
  n_stages = (orig_len-1).bit_length()
  pads = tuple((0, 2**n_stages - orig_len) if i == dim else None for i in range(x.ndim))
  x = x.pad(pads, value=x.dtype.min if descending else x.dtype.max).unflatten(dim, (2,)*n_stages)
  # https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort1.svg
  for stage in range(1, n_stages+1):
    if stage != n_stages:
      # flip so arrows of green boxes point the same way as blue boxes
      crossover_dim = dim + n_stages - stage - 1
      blue_box, green_box = x.split(1, crossover_dim)
      flip_dims = tuple(-i for i in range(1, stage+1+(self.ndim-dim)))
      x = (blue_box.cat(green_box.flip(flip_dims), dim=crossover_dim)).contiguous()
    for substage in range(stage-1, -1, -1):
      partner_dim = dim + n_stages - substage - 1
      x_top, x_bottom = x.split(1, partner_dim)
      x_larger, x_smaller = x_top.maximum(x_bottom), x_top.minimum(x_bottom)
      x = (x_larger.cat(x_smaller, dim=partner_dim) if descending else x_smaller.cat(x_larger, dim=partner_dim)).contiguous()
    if stage != n_stages:
      # flip wires back to undo the crossover
      blue_box, flipped_green_box = x.split(1, crossover_dim)
      x = blue_box.cat(flipped_green_box.flip(flip_dims), dim=crossover_dim)
  x = x.flatten(dim, dim+n_stages-1).shrink_to(self.shape)
  # compute indices for sorted values
  mask = Tensor.ones(orig_len, orig_len, dtype=dtypes.bool, device=self.device).tril().reshape((None, None) + (1,)*(self.ndim-dim-1))
  def compute_counts(t:Tensor): return (mask & (t.unsqueeze(dim) == t.unsqueeze(dim+1))).sum(dim+1)
  count_orig, count_sorted = compute_counts(self), compute_counts(x)
  cond = (self.unsqueeze(dim+1) == x.unsqueeze(dim)) & (count_orig.unsqueeze(dim+1) == count_sorted.unsqueeze(dim))
  idx = Tensor.arange(orig_len, device=self.device).reshape(tuple(orig_len if i == dim else 1 for i in range(x.ndim)))
  idx = (cond * idx.unsqueeze(dim+1)).sum(dim)
  return x, idx

argsort

argsort(dim: int = -1, descending: bool = False) -> Tensor

Returns the indices that sort input tensor along given dimension in given descending order by value.

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

[[3 0 1 2]
 [0 3 2 1]]

Source code in tinygrad/tensor.py

def argsort(self, dim:int=-1, descending:bool=False) -> Tensor:
  """
  Returns the indices that sort input tensor along given `dimension` in given `descending` order by value.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[2, 3, 4, 1], [1, 4, 3, 2]])
  print(t.argsort().numpy())
  ```
  """
  return self.sort(dim, descending)[1]

topk

topk(
    k: int,
    dim: int = -1,
    largest: bool = True,
    sorted_: bool = True,
) -> tuple[Tensor, Tensor]

Computes the top-k elements of the tensor along the specified dim.

Order of indices for equivalent elements is always preserved.

t = Tensor([[0.1, 0.5, 1.2, 3.4, 2.1], [2.2, 1.9, 0.3, 4.5, 0.8]])
print(t.numpy())

[[0.1 0.5 1.2 3.4 2.1]
 [2.2 1.9 0.3 4.5 0.8]]

topk_values, topk_indices = t.topk(2, dim=1)
print(topk_values.numpy())
print(topk_indices.numpy())

[[3.4 2.1]
 [4.5 2.2]]
[[3 4]
 [3 0]]

Source code in tinygrad/tensor.py

def topk(self, k:int, dim:int=-1, largest:bool=True, sorted_:bool=True) -> tuple[Tensor, Tensor]:
  """
  Computes the top-k elements of the tensor along the specified `dim`.

  Order of indices for equivalent elements is always preserved.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[0.1, 0.5, 1.2, 3.4, 2.1], [2.2, 1.9, 0.3, 4.5, 0.8]])
  print(t.numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  topk_values, topk_indices = t.topk(2, dim=1)
  print(topk_values.numpy())
  print(topk_indices.numpy())
  ```
  """
  if not sorted_: raise NotImplementedError("topk with sorted_=False is not supported")
  if k > self.shape[dim:=self._resolve_dim(dim)]: raise ValueError(f"selected index {k=} is out of range")
  x, idx = self.sort(dim, descending=largest)
  topk_shape = tuple(k if i == dim else None for i in range(self.ndim))
  return x.shrink_to(topk_shape), idx.shrink_to(topk_shape)

multinomial

multinomial(
    num_samples: int = 1, replacement: bool = False
) -> Tensor

Returns a tensor with num_samples indices sampled from a multinomial distribution weighted by self.

Note

replacement=False for num_samples > 1 is not supported yet.

Tensor.manual_seed(42)
t = Tensor([1, 2, 3, 4])
print(t.multinomial(20, replacement=True).numpy())

[3 2 3 1 2 3 2 3 0 1 2 3 1 2 2 1 2 2 3 3]

Source code in tinygrad/tensor.py

def multinomial(self:Tensor, num_samples:int = 1, replacement:bool = False) -> Tensor:
  """
  Returns a tensor with `num_samples` indices sampled from a multinomial distribution weighted by `self`.

  NOTE: `replacement=False` for `num_samples > 1` is not supported yet.
  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor([1, 2, 3, 4])
  print(t.multinomial(20, replacement=True).numpy())
  ```
  """
  assert 1 <= self.ndim <= 2 and num_samples > 0, f"{self.ndim=} must be 1 or 2 dim, {num_samples=} must be positive"
  assert replacement or num_samples == 1, "no replacement only supports num_samples = 1"
  weight = self.unsqueeze(0) if self.ndim == 1 else self
  cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
  unif_samples = Tensor.rand(num_samples, cdf.shape[0], 1).to(self.device)
  indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
  return (indices.squeeze(0) if self.ndim == 1 else indices).cast(dtypes.int32)

Neural Network (functional)

linear

linear(
    weight: Tensor,
    bias: Tensor | None = None,
    dtype: DTypeLike | None = None,
) -> Tensor

Applies a linear transformation to self using weight and bias.

See: https://pytorch.org/docs/stable/generated/torch.nn.Linear.html

t = Tensor([[1, 2], [3, 4]])
weight = Tensor([[1, 2], [3, 4]])
bias = Tensor([1, 2])
print(t.linear(weight, bias).numpy())

[[ 8 12]
 [16 24]]

Source code in tinygrad/tensor.py

def linear(self, weight:Tensor, bias:Tensor|None=None, dtype:DTypeLike|None=None) -> Tensor:
  """
  Applies a linear transformation to `self` using `weight` and `bias`.

  See: https://pytorch.org/docs/stable/generated/torch.nn.Linear.html

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[1, 2], [3, 4]])
  weight = Tensor([[1, 2], [3, 4]])
  bias = Tensor([1, 2])
  print(t.linear(weight, bias).numpy())
  ```
  """
  if dtype is not None: return self.cast(dtype).linear(weight.cast(dtype), bias.cast(dtype) if bias is not None else bias)
  x = self.mul(weight) if len(weight.shape) == 1 else self.dot(weight)
  return x.add(bias) if bias is not None else x

sequential

sequential(ll: list[Callable[[Tensor], Tensor]]) -> Tensor

Applies a sequence of functions to self chaining the output of each function to the input of the next.

t = Tensor([1, 2, 3])
print(t.sequential([lambda x: x * 2, lambda x: x + 1]).numpy())

[3 5 7]

Source code in tinygrad/tensor.py

def sequential(self, ll:list[Callable[[Tensor], Tensor]]) -> Tensor:
  """
  Applies a sequence of functions to `self` chaining the output of each function to the input of the next.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([1, 2, 3])
  print(t.sequential([lambda x: x * 2, lambda x: x + 1]).numpy())
  ```
  """
  return functools.reduce(lambda x,f: f(x), ll, self)

layernorm

layernorm(
    axis: int | tuple[int, ...] = -1, eps: float = 1e-05
) -> Tensor

Applies Layer Normalization over a mini-batch of inputs.

• Paper: https://arxiv.org/abs/1607.06450v1

t = Tensor.randn(8, 10, 16) * 2 + 8
print(t.mean().item(), t.std().item())

7.967566967010498 1.9459360837936401

t = t.layernorm()
print(t.mean().item(), t.std().item())

-3.414653448885474e-09 1.0003889799118042

Source code in tinygrad/tensor.py

def layernorm(self, axis:int|tuple[int,...]=-1, eps:float=1e-5) -> Tensor:
  """
  Applies Layer Normalization over a mini-batch of inputs.

  - Paper: https://arxiv.org/abs/1607.06450v1

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.randn(8, 10, 16) * 2 + 8
  print(t.mean().item(), t.std().item())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  t = t.layernorm()
  print(t.mean().item(), t.std().item())
  ```
  """
  y = (self - self.mean(axis, keepdim=True))
  return y.mul((y*y).mean(axis, keepdim=True).add(eps).rsqrt())

batchnorm

batchnorm(
    weight: Tensor | None,
    bias: Tensor | None,
    mean: Tensor,
    invstd: Tensor,
    axis: int | tuple[int, ...] = 1,
) -> Tensor

Applies Batch Normalization over a mini-batch of inputs.

• Paper: https://arxiv.org/abs/1502.03167

t = Tensor.randn(8, 4, 16, 16) * 2 + 8
print(t.mean().item(), t.std().item())

8.022211074829102 2.003443479537964

t = t.batchnorm(None, None, t.mean(axis=(0,2,3)), t.var(axis=(0,2,3)).add(1e-5).rsqrt())
print(t.mean().item(), t.std().item())

9.536140765931123e-08 0.9998152256011963

Source code in tinygrad/tensor.py

def batchnorm(self, weight:Tensor|None, bias:Tensor|None, mean:Tensor, invstd:Tensor, axis:int|tuple[int, ...]=1) -> Tensor:
  """
  Applies Batch Normalization over a mini-batch of inputs.

  - Paper: https://arxiv.org/abs/1502.03167

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor.randn(8, 4, 16, 16) * 2 + 8
  print(t.mean().item(), t.std().item())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  t = t.batchnorm(None, None, t.mean(axis=(0,2,3)), t.var(axis=(0,2,3)).add(1e-5).rsqrt())
  print(t.mean().item(), t.std().item())
  ```
  """
  axis_ = argfix(axis)
  shape = tuple(s if ax in axis_ else 1 for ax, s in enumerate(self.shape))
  x = self - mean.reshape(shape)
  if weight is not None: x = x * weight.reshape(shape)
  ret = x.mul(invstd.reshape(shape) if len(invstd.shape) == len(axis_) else invstd)
  return (ret + bias.reshape(shape)) if bias is not None else ret

dropout

dropout(p=0.5) -> Tensor

Applies dropout to self.

Note

dropout is only applied when Tensor.training is True.

• Paper: https://jmlr.org/papers/v15/srivastava14a.html

Tensor.manual_seed(42)
t = Tensor.randn(2, 2)
with Tensor.train():
  print(t.dropout().numpy())

[[1.2452 0.3412]
 [1.6594 0.6135]]

Source code in tinygrad/tensor.py

def dropout(self, p=0.5) -> Tensor:
  """
  Applies dropout to `self`.

  NOTE: dropout is only applied when `Tensor.training` is `True`.

  - Paper: https://jmlr.org/papers/v15/srivastava14a.html

  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor.randn(2, 2)
  with Tensor.train():
    print(t.dropout().numpy())
  ```
  """
  if not 0 <= p <= 1: raise ValueError(f"{p=} is out of range [0, 1]")
  if not Tensor.training or p == 0: return self
  if p == 1: return self.zeros_like()
  return (Tensor.rand_like(self, requires_grad=False, dtype=dtypes.default_float, contiguous=False) >= p).contiguous().where(self, 0) / (1.0 - p)

one_hot

one_hot(num_classes: int = -1) -> Tensor

Converts self to a one-hot tensor.

num_classes defaults to -1, which means num_classes will be inferred as max(self) + 1.

t = Tensor([0, 1, 3, 3, 4])
print(t.one_hot(5).numpy())

[[1 0 0 0 0]
 [0 1 0 0 0]
 [0 0 0 1 0]
 [0 0 0 1 0]
 [0 0 0 0 1]]

Source code in tinygrad/tensor.py

def one_hot(self, num_classes:int=-1) -> Tensor:
  """
  Converts `self` to a one-hot tensor.

  `num_classes` defaults to -1, which means num_classes will be inferred as max(self) + 1.

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([0, 1, 3, 3, 4])
  print(t.one_hot(5).numpy())
  ```
  """
  if not dtypes.is_int(self.dtype): raise RuntimeError(f"expect integer dtype, getting {self.dtype=}")
  if num_classes == -1: num_classes = int(self.max().item())+1
  return self[..., None]._one_hot_along_dim(num_classes).where(1, 0)

scaled_dot_product_attention

scaled_dot_product_attention(
    key: Tensor,
    value: Tensor,
    attn_mask: Tensor | None = None,
    dropout_p: float = 0.0,
    is_causal: bool = False,
    enable_gqa: bool = False,
) -> Tensor

Computes scaled dot-product attention. self is the query tensor, key is the key tensor, and value is the value tensor.

• Paper: https://arxiv.org/abs/1706.03762v7

q = Tensor.randn(2, 4, 8)
k = Tensor.randn(2, 4, 8)
v = Tensor.randn(2, 4, 8)
print(q.scaled_dot_product_attention(k, v).numpy())

[[[ 1.0179 -0.1569  0.2824 -0.4453 -0.2782 -1.3624 -0.8654  0.0421]
  [ 0.3478 -0.7684  0.9486  1.0626 -0.3239 -0.3466  0.5342 -0.1153]
  [ 0.4184 -0.4923  0.5823  0.3521 -0.2315 -0.6822  0.0829  0.2321]
  [ 0.9673 -0.1948  0.2563 -0.4294 -0.2532 -1.3211 -0.7997  0.0848]]

 [[-0.2036  0.4213  0.8622 -0.9485  0.3362 -0.3886  0.1038  0.3896]
  [-0.0297  1.2289  0.512  -0.3317  0.3861 -0.3695  0.1857  0.4452]
  [-0.4912  0.4212  0.9596 -1.3411  0.4038 -0.5133 -0.2132  0.4209]
  [-0.0036  0.6412  0.8896 -0.6953  0.3864 -0.4349  0.501   0.4907]]]

Source code in tinygrad/tensor.py

def scaled_dot_product_attention(self, key:Tensor, value:Tensor, attn_mask:Tensor|None=None, dropout_p:float=0.0,
                                 is_causal:bool=False, enable_gqa:bool=False) -> Tensor:
  """
  Computes scaled dot-product attention.
  `self` is the query tensor, `key` is the key tensor, and `value` is the value tensor.

  - Paper: https://arxiv.org/abs/1706.03762v7

  ```python exec="true" source="above" session="tensor" result="python"
  q = Tensor.randn(2, 4, 8)
  k = Tensor.randn(2, 4, 8)
  v = Tensor.randn(2, 4, 8)
  print(q.scaled_dot_product_attention(k, v).numpy())
  ```
  """
  if getenv("FLASH_ATTENTION"):
    from extra.thunder.tiny.fa import flash_attention
    return flash_attention(self, key, value, attn_mask=attn_mask, is_causal=is_causal)

  if getenv("HK_FLASH_ATTENTION"):
    from extra.thunder.amd.fa import flash_attention
    return flash_attention(self, key, value, attn_mask=attn_mask, is_causal=is_causal)

  # GQA: https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
  if enable_gqa:
    key = key.repeat_interleave(int(self.shape[-3] // key.shape[-3]), dim=-3)
    value = value.repeat_interleave(int(self.shape[-3] // value.shape[-3]), dim=-3)

  q = self
  qk = q.matmul(key.transpose(-2,-1), dtype=least_upper_dtype(q.dtype, key.dtype, dtypes.float32)) / math.sqrt(q.shape[-1])
  # handle attention mask
  if is_causal:
    if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
    attn_mask = qk.ones_like(requires_grad=False, dtype=dtypes.bool).tril()
  if attn_mask is not None:
    if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
    qk = qk + attn_mask
  return qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value

binary_crossentropy

binary_crossentropy(
    Y: Tensor, reduction: ReductionStr = "mean"
) -> Tensor

Computes the binary cross-entropy loss between self and Y.

See: https://pytorch.org/docs/stable/generated/torch.nn.BCELoss.html

t = Tensor([0.1, 0.9, 0.2])
Y = Tensor([0, 1, 0])
print(t.binary_crossentropy(Y).item())

0.14462155103683472

Source code in tinygrad/tensor.py

def binary_crossentropy(self, Y:Tensor, reduction:ReductionStr="mean") -> Tensor:
  """
  Computes the binary cross-entropy loss between `self` and `Y`.

  See: https://pytorch.org/docs/stable/generated/torch.nn.BCELoss.html

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([0.1, 0.9, 0.2])
  Y = Tensor([0, 1, 0])
  print(t.binary_crossentropy(Y).item())
  ```
  """
  return (-Y*self.log() - (1-Y)*(1-self).log())._do_reduction(reduction)

binary_crossentropy_logits

binary_crossentropy_logits(
    Y: Tensor,
    reduction: ReductionStr = "mean",
    pos_weight: Tensor | None = None,
) -> Tensor

Computes the binary cross-entropy loss between self and Y where self is logits.

See: https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html

t = Tensor([-1, 2, -3])
Y = Tensor([0, 1, 0])
print(t.binary_crossentropy_logits(Y).item())

0.16292566061019897

Source code in tinygrad/tensor.py

def binary_crossentropy_logits(self, Y:Tensor, reduction:ReductionStr="mean", pos_weight:Tensor|None=None) -> Tensor:
  """
  Computes the binary cross-entropy loss between `self` and `Y` where `self` is logits.

  See: https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([-1, 2, -3])
  Y = Tensor([0, 1, 0])
  print(t.binary_crossentropy_logits(Y).item())
  ```
  """
  log_p, log_1_minus_p = self.logsigmoid(), (-self).logsigmoid()
  return (-((1 if pos_weight is None else pos_weight) * Y * log_p + (1-Y) * log_1_minus_p))._do_reduction(reduction)

sparse_categorical_crossentropy

sparse_categorical_crossentropy(
    Y: Tensor,
    ignore_index: int = -1,
    label_smoothing=0.0,
    reduction: ReductionStr = "mean",
) -> Tensor

Computes the sparse categorical cross-entropy loss between self and Y.

Note

self is logits and Y is the target labels. NOTE: unlike PyTorch, this function expects the class axis to be -1

See: https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html

t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.sparse_categorical_crossentropy(Y).item())

0.09391524642705917

Source code in tinygrad/tensor.py

def sparse_categorical_crossentropy(self, Y:Tensor, ignore_index:int=-1, label_smoothing=0.0, reduction:ReductionStr="mean") -> Tensor:
  """
  Computes the sparse categorical cross-entropy loss between `self` and `Y`.

  NOTE: `self` is logits and `Y` is the target labels.
  NOTE: unlike PyTorch, this function expects the class axis to be -1

  See: https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[-1, 2, -3], [1, -2, 3]])
  Y = Tensor([1, 2])
  print(t.sparse_categorical_crossentropy(Y).item())
  ```
  """
  assert 0.0 <= label_smoothing <= 1.0, "label_smoothing must be in [0.0, 1.0]"
  log_probs = self.log_softmax()
  loss_mask = (Y != ignore_index) if ignore_index != -1 else Y.ones_like(dtype=dtypes.bool)
  y = Y.to(self.device).unsqueeze(-1)._one_hot_along_dim(self.shape[-1], dim=-1) * loss_mask.unsqueeze(-1)
  smoothing = label_smoothing * (log_probs.mean(-1) * loss_mask)
  unreduced = ((1 - label_smoothing) * (log_probs * y).sum(-1) + smoothing)
  return -unreduced.sum() / loss_mask.sum() if reduction == "mean" else -unreduced._do_reduction(reduction)

cross_entropy

cross_entropy(
    Y: Tensor,
    reduction: ReductionStr = "mean",
    label_smoothing: float = 0.0,
) -> Tensor

Computes the cross entropy loss between input logits and target.

Note

self are logits and Y are the target labels or class probabilities.

See: https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html

t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.cross_entropy(Y).item())

0.09391524642705917

t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.cross_entropy(Y, reduction='none').numpy())

[0.055  0.1328]

Source code in tinygrad/tensor.py

def cross_entropy(self, Y:Tensor, reduction:ReductionStr="mean", label_smoothing:float=0.0) -> Tensor:
  """
  Computes the cross entropy loss between input logits and target.

  NOTE: `self` are logits and `Y` are the target labels or class probabilities.

  See: https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[-1, 2, -3], [1, -2, 3]])
  Y = Tensor([1, 2])
  print(t.cross_entropy(Y).item())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[-1, 2, -3], [1, -2, 3]])
  Y = Tensor([1, 2])
  print(t.cross_entropy(Y, reduction='none').numpy())
  ```
  """
  assert 0.0 <= label_smoothing <= 1.0, "label_smoothing must be in [0.0, 1.0]"
  classes_dim = 0 if self.ndim == 1 else 1
  if self.shape != Y.shape:
    if self.max(classes_dim).shape != Y.shape: raise RuntimeError(f"shape mismatch: {self.shape=}, {Y.shape=}")
    Y = Y.unsqueeze(classes_dim)._one_hot_along_dim(num_classes=self.shape[classes_dim], dim=classes_dim)
  Y = (1 - label_smoothing)*Y + label_smoothing / int(Y.shape[classes_dim])
  return -self.log_softmax(classes_dim).mul(Y).sum(classes_dim)._do_reduction(reduction)

nll_loss

nll_loss(
    Y: Tensor,
    weight: Tensor | None = None,
    ignore_index: int | None = None,
    reduction: ReductionStr = "mean",
) -> Tensor

Computes the negative log likelihood loss between log-probabilities and target labels.

Note

self is log-probabilities and Y is the Y labels or class probabilities.

See: https://pytorch.org/docs/stable/generated/torch.nn.functional.nll_loss.html

t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.log_softmax().nll_loss(Y).item())

0.09391524642705917

t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.log_softmax().nll_loss(Y, reduction='none').numpy())

[0.055  0.1328]

Source code in tinygrad/tensor.py

def nll_loss(self, Y:Tensor, weight:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean") -> Tensor:
  """
  Computes the negative log likelihood loss between log-probabilities and target labels.

  NOTE: `self` is log-probabilities and `Y` is the Y labels or class probabilities.

  See: https://pytorch.org/docs/stable/generated/torch.nn.functional.nll_loss.html

  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[-1, 2, -3], [1, -2, 3]])
  Y = Tensor([1, 2])
  print(t.log_softmax().nll_loss(Y).item())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  t = Tensor([[-1, 2, -3], [1, -2, 3]])
  Y = Tensor([1, 2])
  print(t.log_softmax().nll_loss(Y, reduction='none').numpy())
  ```
  """
  weight = Y.ones_like(requires_grad=False) if weight is None else weight[Y]
  masked_weight = weight if ignore_index is None else weight * (Y != ignore_index)
  nll = -self.gather(1, Y.unsqueeze(1)).squeeze(1) * masked_weight
  return nll.sum() / masked_weight.sum() if reduction == "mean" else nll._do_reduction(reduction)

Linear Algebra

qr

qr() -> tuple[Tensor, Tensor]

Source code in tinygrad/tensor.py

def qr(self) -> tuple[Tensor, Tensor]:
  assert self.ndim > 1, f"expected two or more dimensions, got {self.ndim}"
  b_shape, m, n = self.shape[:-2], int(self.shape[-2]), int(self.shape[-1])
  R = self.clone()
  Q = Tensor.eye(m, dtype=self.dtype).reshape((1,) * len(b_shape) + (m, m)).expand(b_shape + (m, m))
  for i in range(min(m, n)):
    x = R[..., i:m, i]
    norm = x.square().sum(-1).sqrt()
    s = (x[..., 0] != 0).where(-x[..., 0].sign(), -1)
    u1 = x[..., 0] - s * norm
    w = x.unsqueeze(-1) / (norm != 0).where(u1, 1).reshape(b_shape + (1, 1))
    w[..., 0, 0] = 1
    tau = (-s * u1 / (norm != 0).where(norm, 1)).reshape(b_shape + (1, 1))
    tau = (norm != 0).reshape(b_shape + (1, 1)).where(tau, 0)
    R[..., i:m, :] = R[..., i:m, :] - (w * tau) @ (w.transpose(-2, -1) @ R[..., i:m, :])
    Q[..., :, i:m] = Q[..., :, i:m] - (Q[..., :, i:m] @ w) @ (tau * w).transpose(-2, -1)
  return Q,R

svd

svd(full_matrices=True) -> tuple[Tensor, Tensor, Tensor]

Source code in tinygrad/tensor.py

def svd(self, full_matrices = True) -> tuple[Tensor, Tensor, Tensor]:
  #partial implementation of https://www.netlib.org/lapack/lawnspdf/lawn169.pdf , pg 26
  assert self.ndim > 1, f"expected two or more dimensions, got {self.ndim}"
  b_shape, m, n = self.shape[:-2], int(self.shape[-2]), int(self.shape[-1])
  #preprocess the matrix
  Q, R = (self.qr() if m >= n else self.transpose(-2, -1).qr())
  num, q_num = min(m, n), max(m, n)
  # TODO: codegen infinite loop without contiguous
  U = R.shrink(tuple([None] * len(b_shape) + [(0, num), (0, num)])).contiguous()
  V = Tensor.eye(num, dtype=self.dtype).reshape((1,) * len(b_shape) + (num, num)).expand(b_shape + (num, num)).contiguous()
  #prepare round robin pairing
  permute, inverse_permute = Tensor.arange(0, num, dtype=dtypes.int), Tensor.zeros(num, dtype=dtypes.int)
  permute[num//2:num] = permute[num//2:num].flip(0)
  inverse_permute[permute] = Tensor.arange(num, dtype=dtypes.int)
  def one_round_jacobi(U, V,permute,inverse_permute):
    #pair all the columns
    V_permuted, runoff_V = (V[..., permute].split(num - 1, -1)) if num % 2 == 1 else (V[..., permute], None)
    V_left, V_right = V_permuted.split(num//2, -1)
    U_permuted, runoff_U = (U[..., permute].split(num - 1, -1)) if num % 2 == 1 else (U[..., permute], None)
    U_left, U_right = U_permuted.split(num//2, -1)
    #compute the jacobi rotations for each pairing
    gamma = (U_left * U_right).sum(-2).reshape(b_shape + (1, num//2))
    alpha, beta = U_permuted.square().sum(-2).unsqueeze(-2).split(num//2, -1)
    rot = gamma != 0
    tau = (beta - alpha) / (2 * rot.where(gamma, 1))
    t = (tau != 0).where(tau.sign(), 1) / (tau.abs() + (1 + tau.square()).sqrt())
    t = rot.where(t, 0)
    c = 1 / (1 + t.square()).sqrt()
    s = c * t
    #apply the rotations
    U_left, U_right = c * U_left - s * U_right, s * U_left + c * U_right
    U = U_left.cat(U_right.cat(runoff_U, dim = -1) if num % 2 == 1 else U_right, dim = -1)[..., inverse_permute]
    V_left, V_right = c * V_left - s * V_right, s * V_left + c * V_right
    V = V_left.cat(V_right.cat(runoff_V, dim = -1) if num % 2 == 1 else V_right, dim = -1)[..., inverse_permute]
    #prepare the next round robin pairings
    if num % 2 == 1: permute = ((permute - 1) % num)
    else: permute = permute[0].reshape(1).cat(((permute[1:num] - 2) % (num - 1)) + 1)
    inverse_permute = inverse_permute.scatter(0,permute,Tensor.arange(num,dtype=dtypes.int32))
    return U, V, permute, inverse_permute
  max_iterations, iterations_per_round = 1, int(num * math.log2(num) * 2 + 2)#sorta heuristic, most use num*log2(num)
  for _ in range(max_iterations * iterations_per_round): U, V, permute, inverse_permute = one_round_jacobi(U, V, permute, inverse_permute)
  #extract singular values and sort. construct U from Q
  S, indices = U.square().sum(-2).sqrt().sort(dim = -1, descending=True)
  new_indices = indices.reshape(b_shape + (1, num)).expand(b_shape + (num, num))
  U = U.gather(-1, new_indices) / (S != 0).where(S, 1).unsqueeze(-2)
  V = V.gather(-1, new_indices)

  padded_u = Tensor.eye(q_num, dtype=U.dtype).reshape((1,) * len(b_shape) + (q_num, q_num)).expand(b_shape + (q_num, q_num))
  padded_u[..., 0:num, 0:num] = U
  U = Q @ padded_u
  if not full_matrices: U, V = U[..., 0:num], V[..., 0:num]
  return (U, S, V.transpose(-2,-1)) if m >= n else (V, S, U.transpose(-2, -1))