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
) -> Self

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/mixin/op.py

def min(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Self:
  """
  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
) -> Self

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/mixin/reduce.py

def any(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Self:
  """
  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
) -> Self

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/mixin/reduce.py

def all(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Self:
  """
  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().prod(axis, keepdim)

isclose

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

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/mixin/elementwise.py

def isclose(self, other, rtol:float=1e-05, atol:float=1e-08, equal_nan=False) -> Self:
  """
  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())
  ```
  """
  other = self.ufix(other)
  is_finite_close = self.isfinite() & other.isfinite() & ((self - other).abs() <= atol + rtol * other.abs())
  is_infinite_close = (self.isinf() | other.isinf()) & self.eq(other)
  is_nan_close = (self.isnan() & other.isnan()) & equal_nan
  return is_finite_close | is_infinite_close | is_nan_close

allclose

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

Check if all self and other are close.

Source code in tinygrad/mixin/op.py

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

mean

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

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/mixin/op.py

def mean(self, axis:int|Sequence[int]|None=None, keepdim=False) -> Self:
  """
  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)  # type: ignore[arg-type]

var

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

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/mixin/op.py

def var(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> Self:
  """
  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)])
  reduced = squares.sum(axis=axis, keepdim=keepdim)
  denominator = reduced.const_like(n) - correction  # type: ignore[arg-type]
  # TODO: remove relu?
  return reduced.div(denominator.relu())

var_mean

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

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/mixin/op.py

def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Self, Self]:
  """
  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,
) -> Self

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/mixin/op.py

def std(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> Self:
  """
  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[Self, Self]

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/mixin/op.py

def std_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Self, Self]:
  """
  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) -> Self

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/mixin/op.py

def softmax(self, axis=-1, dtype:DTypeLike|None=None) -> Self:
  """
  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 * ss.reciprocal()

log_softmax

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

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/mixin/op.py

def log_softmax(self, axis=-1, dtype:DTypeLike|None=None) -> Self:
  """
  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) -> Self

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/mixin/op.py

def logsumexp(self, axis=None, keepdim=False) -> Self:
  """
  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) -> Self

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/mixin/op.py

def logcumsumexp(self, axis=0) -> Self:
  """
  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 = type(self).ones(last_dim_size, last_dim_size, buffer=False).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) -> Self

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/mixin/op.py

def argmax(self, axis=None, keepdim=False) -> Self:
  """
  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.eq(self.max(axis=axis, keepdim=True))
  idx = m * type(self).arange(self.shape[axis], 0, -1).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) -> Self

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/mixin/op.py

def argmin(self, axis=None, keepdim=False) -> Self:
  """
  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,
) -> Self

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/mixin/op.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) -> Self:
  """
  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))
  s_ = stride if stride is not None else k_
  def pool(x:Self, padding_:Sequence[int]) -> Self:
    return x._pad_constant(((0,0),)*(x.ndim-len(k_)) + flat_to_grouped(padding_), 0.0)._pool(k_, s_, dilation)
  reg_pads = resolve_pool_pads(padding, len(k_))
  pads = self._apply_ceil_mode(reg_pads, k_, s_, dilation) if ceil_mode else reg_pads
  if not count_include_pad:
    return pool(self, pads).sum(axis) / pool(self.const_like(1), pads).sum(axis)
  if not ceil_mode: return pool(self, pads).mean(axis)
  return pool(self, pads).sum(axis) / pool(self._pad_constant(((0,0),)*(self.ndim-len(k_)) + flat_to_grouped(reg_pads), 0.0).const_like(1),
                                            tuple(cp-rp for cp,rp in zip(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,
) -> Self | tuple[Self, Self]

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/mixin/op.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) -> Self | tuple[Self, Self]:
  """
  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))
  s_ = stride if stride is not None else k_
  pads = resolve_pool_pads(padding, len(k_))
  if ceil_mode: pads = self._apply_ceil_mode(pads, k_, s_, dilation)
  pooled = self._pad_constant(((0,0),)*(self.ndim-len(k_)) + flat_to_grouped(pads), self.dtype.min)._pool(k_, s_, dilation)
  if not return_indices: return pooled.max(axis)
  spatial_sz = int(prod(spatial_shape := self.shape[-len(k_):]))
  idx = type(self).arange(spatial_sz, 0, -1).reshape(spatial_shape)
  m = pooled.eq(pooled.max(axis, keepdim=True))
  idx = m * idx._pad_constant(((0,0),)*(idx.ndim-len(k_)) + flat_to_grouped(pads), idx.dtype.min)._pool(k_, s_, dilation)
  return pooled.max(axis), spatial_sz - idx.max(axis)

max_unpool2d

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

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/mixin/op.py

def max_unpool2d(self, indices:Self, kernel_size:tuple[int, ...]=(2,2), stride=None, dilation=1, padding:int|tuple[int, ...]=0,
                 output_size=None) -> Self:
  """
  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(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: Self,
    bias: Self | None = None,
    groups=1,
    stride=1,
    dilation=1,
    padding: int | Sequence[int] = 0,
    dtype: DTypeLike | None = None,
) -> Self

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/mixin/op.py

def conv2d(self, weight:Self, bias:Self|None=None, groups=1, stride=1, dilation=1, padding:int|Sequence[int]=0,
           dtype:DTypeLike|None=None) -> Self:
  """
  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)
  if WINO and all(x == 3 for x in weight.shape[2:]) and stride == dilation == 1: return self._conv2d_winograd(weight, bias, groups, padding, dtype)
  (bs,cin_), (cout,cin), HW = self.shape[:2], weight.shape[:2], weight.shape[2:]
  padding_ = 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, possibly negative — _pad_constant handles the shrink)
  x = self._pad_constant(((0,0),)*(self.ndim-len(HW)) + flat_to_grouped(padding_), 0.0)._pool(HW, stride, dilation)
  rcout, oyx = cout//groups, x.shape[2:-len(HW)]
  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)))

conv_transpose2d

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

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/mixin/op.py

def conv_transpose2d(self, weight:Self, bias:Self|None=None, groups=1, stride=1, dilation=1, padding=0, output_padding=0) -> Self:
  """
  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(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: Self, dtype: DTypeLike | None = None) -> Self

Performs dot product between two tensors. If w is 1-D, it's a sum product over the last axis of self and w. If w is N-D with N>=2, it's a sum product over the last axis of self and the second-to-last axis of w.

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

a = Tensor([1, 2, 3])
b = Tensor([1, 1, 0])
print(a.dot(b).numpy())

3

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

[[19 22]
 [43 50]]

Source code in tinygrad/mixin/op.py

def dot(self, w:Self, dtype:DTypeLike|None=None) -> Self:
  """
  Performs dot product between two tensors.
  If `w` is 1-D, it's a sum product over the last axis of `self` and `w`.
  If `w` is N-D with N>=2, it's a sum product over the last axis of `self` and the second-to-last axis of `w`.

  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])
  b = Tensor([1, 1, 0])
  print(a.dot(b).numpy())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  a = Tensor([[1, 2], [3, 4]])
  b = Tensor([[5, 6], [7, 8]])
  print(a.dot(b).numpy())
  ```
  """
  if IMAGE: return self.image_dot(w, dtype)
  x, dx, dw = self, self.ndim, w.ndim
  if not (dx > 0 and dw > 0): raise RuntimeError(f"both tensors need to be at least 1D, got {dx}D and {dw}D")
  if x.shape[-1] != w.shape[axis_w:=-min(w.ndim,2)]: raise RuntimeError(f"cannot dot {x.shape} and {w.shape}")
  x = x.reshape(*x.shape[0:-1], *[1]*min(dx-1, dw-1, 1), x.shape[-1])
  w = w.reshape(*w.shape[0:-2], *[1]*min(dx-1, dw-1, 1), *w.shape[axis_w:]).transpose(-1, axis_w)
  return (x*w).sum(-1, dtype=dtype).cast(least_upper_dtype(x.dtype, w.dtype) if dtype is None else to_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/op.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 classmethod

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

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/mixin/op.py

@classmethod
def einsum(cls, formula:str, *operands:Self|Sequence[Self], dtype:DTypeLike|None=None) -> Self:
  """
  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), 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 xs[0].uprod(*xs[1:]).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) -> Self

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/mixin/op.py

def cumsum(self, axis:int=0) -> Self:
  """
  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) -> Self

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/mixin/op.py

def cumprod(self, axis:int) -> Self:
  """
  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[Self, Self]

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/mixin/op.py

def cummax(self, axis:int=0) -> tuple[Self, Self]:
  """
  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), type(self).zeros(self.shape, dtype=dtypes.int32, buffer=False)
  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).eq(values_t.unsqueeze(-2)) * type(self).ones(n, n, buffer=False).triu()
  idx = (-(match * type(self).arange(n, 0, -1).reshape(n, 1)).max(-2) + n).cast(dtypes.int32)
  return values, idx.transpose(-1, axis)

cummin

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

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/mixin/op.py

def cummin(self, axis:int=0) -> tuple[Self, Self]:
  """
  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) -> Self

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/mixin/op.py

def triu(self, diagonal:sint=0) -> Self:
  """
  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 self._tri(self.shape[-2], self.shape[-1], diagonal).where(self, self.const_like(0))

tril

tril(diagonal: sint = 0) -> Self

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/mixin/op.py

def tril(self, diagonal:sint=0) -> Self:
  """
  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 self._tri(self.shape[-2], self.shape[-1], diagonal+1).where(self.const_like(0), self)

interpolate

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

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/mixin/op.py

def interpolate(self, size:tuple[int, ...], mode:str="linear", align_corners:bool=False) -> Self:
  """
  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 = type(self).arange(size[i], dtype=dtypes.float32), [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: Self,
    src: Self | PyConst,
    reduce: Literal["multiply", "add"] | None = None,
) -> Self

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/mixin/op.py

def scatter(self, dim:int, index:Self, src:Self|PyConst, reduce:Literal['multiply', 'add']|None=None) -> Self:
  """
  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 isinstance(src, (int, float, bool)): src = type(self).full(index.shape, src, dtype=self.dtype, buffer=False)
  elif reduce: raise TypeError("non-scalar src is not supported with reduce arg. use scatter_reduce")
  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 self._masked_merge(src, mask, (-1,))

scatter_reduce

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

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/mixin/op.py

def scatter_reduce(self, dim:int, index:Self, src:Self, reduce:Literal["sum", "prod", "mean", "amax", "amin"],
                   include_self:bool=True) -> Self:
  """
  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:Self|PyConst, b:Self|PyConst) -> Self: 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, size: int | None = None, fill_value: ConstType = 0
)

Selects elements from self based on the boolean mask.

With size=None (default), output length equals the number of True values (not jittable). With size=N, output length is N, padded with fill_value or truncated (jittable).

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]

print(t.masked_select(mask, size=6, fill_value=-1).numpy())

[ 0  2  4  8 -1 -1]

Source code in tinygrad/mixin/op.py

def masked_select(self, mask, size:int|None=None, fill_value:ConstType=0):
  """
  Selects elements from `self` based on the boolean `mask`.

  With `size=None` (default), output length equals the number of `True` values (not jittable).
  With `size=N`, output length is `N`, padded with `fill_value` or truncated (jittable).

  ```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())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.masked_select(mask, size=6, fill_value=-1).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()
  if size is None:
    counts = type(self).zeros(mask_cumsum[-1].item() if mask.numel() else 0, dtype=dtypes.int32, buffer=False)
    return x[counts.scatter(0, mask_cumsum, 1, reduce='add').cumsum()]
  counts = type(self).zeros(size, dtype=dtypes.int32, buffer=False).scatter(0, mask_cumsum, 1, reduce='add')
  return (type(self).arange(size) < mask.sum()).where(x[counts.cumsum()], fill_value).cast(self.dtype)

masked_fill

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

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/mixin/elementwise.py

def masked_fill(self, mask:Self, value:Self|PyConst) -> Self:
  """
  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(
    size: int | None = None, fill_value: ConstType = 0
) -> Self

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

With size=None (default), output shape is (n_nonzero, ndim) (not jittable). With size=N, output shape is (N, ndim), padded with fill_value or truncated (jittable).

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]]

print(t.nonzero(size=3, fill_value=-1).numpy())

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

Source code in tinygrad/mixin/op.py

def nonzero(self, size:int|None=None, fill_value:ConstType=0) -> Self:
  """
  Returns the indices of the elements that are non-zero.

  With `size=None` (default), output shape is `(n_nonzero, ndim)` (not jittable).
  With `size=N`, output shape is `(N, ndim)`, padded with `fill_value` or truncated (jittable).

  ```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())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  print(t.nonzero(size=3, fill_value=-1).numpy())
  ```
  """
  if self.ndim == 0:
    return type(self).zeros(size if size is not None else int(self.ne(0).item()), 0, dtype=dtypes.int32, device=self.device)
  mask = self.ne(0).flatten()
  indices = type(self).stack(*[type(self).arange(s).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),
                               size=size*self.ndim if size is not None else None, fill_value=fill_value).reshape(-1, self.ndim)

sort

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

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/mixin/op.py

def sort(self, dim:int=-1, descending:bool=False) -> tuple[Self, Self]:
  """
  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.const_like(0).cast(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_constant(pads, 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 = type(self).ones(orig_len, orig_len, dtype=dtypes.bool, buffer=False).tril()
  mask = mask.reshape((None, None) + (1,)*(self.ndim-dim-1))
  def compute_counts(t:Self): return (mask & t.unsqueeze(dim).eq(t.unsqueeze(dim+1))).sum(dim+1)
  count_orig, count_sorted = compute_counts(self), compute_counts(x)
  cond = self.unsqueeze(dim+1).eq(x.unsqueeze(dim)) & count_orig.unsqueeze(dim+1).eq(count_sorted.unsqueeze(dim))
  idx = type(self).arange(orig_len).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) -> Self

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/mixin/op.py

def argsort(self, dim:int=-1, descending:bool=False) -> Self:
  """
  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[Self, Self]

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/mixin/op.py

def topk(self, k:int, dim:int=-1, largest:bool=True, sorted_:bool=True) -> tuple[Self, Self]:
  """
  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
) -> Self

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

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]

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

[2 1 3]

Source code in tinygrad/mixin/rand.py

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

  ```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())
  ```
  ```python exec="true" source="above" session="tensor" result="python"
  Tensor.manual_seed(42)
  t = Tensor([1, 2, 3, 4])
  print(t.multinomial(3, replacement=False).numpy())
  ```
  """
  assert 1 <= self.ndim <= 2 and num_samples > 0, f"{self.ndim=} must be 1 or 2 dim, {num_samples=} must be positive"
  weight = self.unsqueeze(0) if self.ndim == 1 else self
  assert replacement or num_samples <= weight.shape[1], "no replacement samples must not exceed population size"
  if replacement or num_samples == 1:
    cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1)
    unif_samples = type(self).rand(num_samples, cdf.shape[0], 1).to(self.device)  # type: ignore[attr-defined]
    indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0))
  else:
    # Efraimidis-Spirakis
    indices = (weight.rand_like(dtype=dtypes.float32).log2() / weight).topk(num_samples, dim=1)[1]
  return (indices.squeeze(0) if self.ndim == 1 else indices).cast(dtypes.int32)

Neural Network (functional)

linear

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

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/mixin/op.py

def linear(self, weight:Self, bias:Self|None=None, dtype:DTypeLike|None=None) -> Self:
  """
  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:
    dt = to_dtype(dtype)
    return self.cast(dt).linear(weight.cast(dt), bias.cast(dt) 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[[Self], Self]]) -> Self

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/mixin/op.py

def sequential(self, ll:list[Callable[[Self], Self]]) -> Self:
  """
  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
) -> Self

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.882682800292969 1.9711600542068481

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

-1.558376183652399e-08 1.0003888607025146

Source code in tinygrad/mixin/op.py

def layernorm(self, axis:int|tuple[int,...]=-1, eps:float=1e-5) -> Self:
  """
  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: Self | None,
    bias: Self | None,
    mean: Self,
    invstd: Self,
    axis: int | tuple[int, ...] = 1,
) -> Self

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())

7.988671779632568 1.9851549863815308

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())

-1.2356950946923462e-06 0.9998151063919067

Source code in tinygrad/mixin/op.py

def batchnorm(self, weight:Self|None, bias:Self|None, mean:Self, invstd:Self, axis:int|tuple[int, ...]=1) -> Self:
  """
  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) -> Self

Applies dropout to self.

Note

dropout is only applied when TRAINING is set (e.g. inside Context(TRAINING=1)).

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

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

[[1.2452 0.3412]
 [1.6594 0.6135]]

Source code in tinygrad/mixin/rand.py

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

  NOTE: dropout is only applied when `TRAINING` is set (e.g. inside `Context(TRAINING=1)`).

  - 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 Context(TRAINING=1):
    print(t.dropout().numpy())
  ```
  """
  if not 0 <= p <= 1: raise ValueError(f"{p=} is out of range [0, 1]")
  if not TRAINING or p == 0: return self
  if p == 1: return self.const_like(0)
  return (self.rand_like(dtype=dtypes.default_float, contiguous=False) >= p).contiguous().where(self, 0) / (1.0 - p)

one_hot

one_hot(num_classes: int) -> Self

Converts self to a one-hot tensor.

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/mixin/op.py

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

  ```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 < 0: raise ValueError(f"num_classes must be non-negative, got {num_classes}")
  return self[..., None]._one_hot_along_dim(num_classes).where(1, 0)

scaled_dot_product_attention

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

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/mixin/rand.py

def scaled_dot_product_attention(self, key:Self, value:Self, attn_mask:Self|None=None, dropout_p:float=0.0,
                                 is_causal:bool=False, enable_gqa:bool=False) -> Self:
  """
  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())
  ```
  """
  # 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.const_like(1).cast(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: Self, reduction: ReductionStr = "mean"
) -> Self

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/mixin/op.py

def binary_crossentropy(self, Y:Self, reduction:ReductionStr="mean") -> Self:
  """
  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: Self,
    reduction: ReductionStr = "mean",
    pos_weight: Self | None = None,
) -> Self

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/mixin/op.py

def binary_crossentropy_logits(self, Y:Self, reduction:ReductionStr="mean", pos_weight:Self|None=None) -> Self:
  """
  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: Self,
    ignore_index: int = -1,
    label_smoothing=0.0,
    reduction: ReductionStr = "mean",
) -> Self

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/mixin/op.py

def sparse_categorical_crossentropy(self, Y:Self, ignore_index:int=-1, label_smoothing=0.0, reduction:ReductionStr="mean") -> Self:
  """
  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]"
  if Y.device is not None and self.device is not None and Y.device != self.device:
    raise RuntimeError(f"expected Y and self on the same device, {Y.device=}, {self.device=}")
  log_probs = self.log_softmax()
  loss_mask = Y.ne(ignore_index) if ignore_index != -1 else Y.const_like(1).cast(dtypes.bool)
  y = Y.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: Self,
    reduction: ReductionStr = "mean",
    label_smoothing: float = 0.0,
) -> Self

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/mixin/op.py

def cross_entropy(self, Y:Self, reduction:ReductionStr="mean", label_smoothing:float=0.0) -> Self:
  """
  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: Self,
    weight: Self | None = None,
    ignore_index: int | None = None,
    reduction: ReductionStr = "mean",
) -> Self

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/mixin/op.py

def nll_loss(self, Y:Self, weight:Self|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean") -> Self:
  """
  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.const_like(1) if weight is None else weight.gather(0, Y.flatten()).reshape(Y.shape)
  masked_weight = weight if ignore_index is None else weight * Y.ne(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[Self, Self]

Source code in tinygrad/mixin/op.py

def qr(self) -> tuple[Self, Self]:
  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, Q = self, type(self).eye(m, dtype=self.dtype).expand(b_shape + (m, m))
  idx = type(self).arange(m)
  for i in range(min(m, n)):
    # full-length Householder reflector v with zeros above row i; w = tau*v is the rank-1 update factor
    at_i, x = idx.eq(i), (idx >= i).where(R[..., :, i], 0)
    norm = x.square().sum(-1, keepdim=True).sqrt()
    x0 = at_i.where(x, 0).sum(-1, keepdim=True)
    sgn, active = x0.ne(0).where(x0.sign(), 1), norm.ne(0)
    u0 = x0 + sgn * norm
    v = (at_i.where(u0, x) / active.where(u0, 1)).unsqueeze(-1)
    w = active.where(sgn * u0 / active.where(norm, 1), 0).unsqueeze(-1) * v
    R = R - w @ (v.transpose(-2, -1) @ R)
    Q = Q - (Q @ v) @ w.transpose(-2, -1)
  return Q, R

svd

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

Source code in tinygrad/mixin/op.py

def svd(self, full_matrices = True) -> tuple[Self, Self, Self]:
  #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 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[..., :num, :num].contiguous()
  V = type(self).eye(num, dtype=self.dtype).expand(b_shape + (num, num)).contiguous()
  #prepare round robin pairing: identity on first half, reversed on second half
  permute = type(self).arange(num//2, dtype=dtypes.int).cat(
              type(self).arange(num//2, num, dtype=dtypes.int).flip(0))
  cols, h = type(self).arange(num, dtype=dtypes.int), num // 2
  eye_num = type(self).eye(num, dtype=self.dtype).expand(b_shape + (num, num))
  def one_round_jacobi(U, V, permute):
    # permutation matrix P with P[a,b] = (a == permute[b]); first 2h columns are paired-column selectors
    P = cols.unsqueeze(1).eq(permute.unsqueeze(0)).cast(U.dtype)
    P_pair = P[..., :2*h]  # drops the runoff column for odd num
    # extract paired columns to compute Jacobi rotation params
    U_pair = U @ P_pair
    U_left, U_right = U_pair.split(h, -1)
    gamma = (U_left * U_right).sum(-2).reshape(b_shape + (1, h))
    alpha, beta = U_pair.square().sum(-2).unsqueeze(-2).split(h, -1)
    rot = gamma.ne(0)
    tau = (beta - alpha) / (2 * rot.where(gamma, 1))
    t = tau.ne(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
    # build rotation matrix R: identity + sum over pairs of 2x2 rotation deltas at (i_k, j_k) positions
    Mi, Mj = P_pair.transpose(-2, -1).split(h, -2)  # paired-column selectors, each shape (h, num)
    Mi_a, Mi_b = Mi.unsqueeze(-1), Mi.unsqueeze(-2)
    Mj_a, Mj_b = Mj.unsqueeze(-1), Mj.unsqueeze(-2)
    cc, ss = (c - 1).reshape(b_shape + (h, 1, 1)), s.reshape(b_shape + (h, 1, 1))
    R = eye_num + (cc * (Mi_a * Mi_b + Mj_a * Mj_b) + ss * (Mi_a * Mj_b - Mj_a * Mi_b)).sum(-3)
    U, V = U @ R, V @ R
    #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)
    return U, V, permute
  # classical Jacobi converges in ~4 sweeps; one full sweep is (num-1) rounds for even num
  for _ in range(4 * num): U, V, permute = one_round_jacobi(U, V, 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.unsqueeze(-2).expand(b_shape + (num, num))
  U = U.gather(-1, new_indices) / S.ne(0).where(S, 1).unsqueeze(-2)
  V = V.gather(-1, new_indices)
  # place U into the top-left num×num block of a q_num×q_num identity matrix
  pad_arg = (None,) * len(b_shape) + ((0, q_num - num), (0, q_num - num))
  eye_q = type(self).eye(q_num, dtype=U.dtype).expand(b_shape + (q_num, q_num))
  eye_n = type(self).eye(num, dtype=U.dtype).expand(b_shape + (num, num)).pad(pad_arg)
  U = Q @ (U.pad(pad_arg) + eye_q - eye_n)
  if not full_matrices: U = U[..., 0:num]
  return (U, S, V.transpose(-2, -1)) if m >= n else (V, S, U.transpose(-2, -1))