From 8091661df34bbc12df585f124aca50e0e007ff41 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Sun, 15 Feb 2026 15:18:37 +0800 Subject: [PATCH] more more to mixins (#14761) --- tinygrad/mixin/math.py | 70 ++++++++++++++++++++++++++++ tinygrad/mixin/movement.py | 26 ++++++++++- tinygrad/tensor.py | 94 -------------------------------------- 3 files changed, 95 insertions(+), 95 deletions(-) diff --git a/tinygrad/mixin/math.py b/tinygrad/mixin/math.py index f38d6cc9d0..8c0a58c3d6 100644 --- a/tinygrad/mixin/math.py +++ b/tinygrad/mixin/math.py @@ -514,3 +514,73 @@ class MathMixin: ``` """ return self.sqrt().reciprocal() + + def log(self) -> Self: + """ + Computes the natural logarithm element-wise. + + See: https://en.wikipedia.org/wiki/Logarithm + + ```python exec="true" source="above" session="tensor" result="python" + print(Tensor([1., 2., 4., 8.]).log().numpy()) + ``` + """ + return self.log2()*math.log(2) + + def log10(self) -> Self: + """ + Computes the base-10 logarithm element-wise. + + See: https://en.wikipedia.org/wiki/Logarithm + + ```python exec="true" source="above" session="tensor" result="python" + print(Tensor([1., 2., 4., 8.]).log10().numpy()) + ``` + """ + return self.log2()*math.log10(2) + + def atanh(self) -> Self: + """ + Applies the Inverse Hyperbolic Tangent (atanh) function element-wise. + + - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#atanh + + ```python exec="true" source="above" session="tensor" result="python" + print(Tensor([-0.9, -0.6, -0.3, 0., 0.3, 0.6, 0.9]).atanh().numpy()) + ``` + """ + return ((1 + self)/(1 - self)).log() / 2 + + def asinh(self) -> Self: + """ + Applies the Inverse Hyperbolic Sine (asinh) function element-wise. + + - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#asinh + + ```python exec="true" source="above" session="tensor" result="python" + print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).asinh().numpy()) + ``` + """ + return (self + (self.square() + 1).sqrt()).log() + + def acosh(self) -> Self: + """ + Applies the Inverse Hyperbolic Cosine (acosh) function element-wise. + + - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#acosh + + ```python exec="true" source="above" session="tensor" result="python" + print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).acosh().numpy()) + ``` + """ + return (self + (self.square() - 1).sqrt()).log() + + def round(self) -> Self: + """ + Rounds the tensor element-wise with rounding half to even. + + ```python exec="true" source="above" session="tensor" result="python" + print(Tensor([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]).round().numpy()) + ``` + """ + return ((self > 0).eq((b := self.trunc() / 2.0).trunc().eq(b))).where((self - 0.5).ceil(), (self + 0.5).floor()) diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py index 0a901716f0..aafb9d67b0 100644 --- a/tinygrad/mixin/movement.py +++ b/tinygrad/mixin/movement.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, Self from tinygrad.uop import Ops -from tinygrad.helpers import prod, argfix, flatten, dedup, make_tuple, ceildiv +from tinygrad.helpers import prod, argfix, argsort, flatten, dedup, make_tuple, ceildiv from tinygrad.uop.ops import resolve, smax if TYPE_CHECKING: @@ -373,3 +373,27 @@ class MovementMixin: x = x.shrink_to(noop + flatten((k, o, 1) for k, o in zip(k_, o_))).reshape(noop + flatten((k, o) for k, o in zip(k_, o_))) # permute to move reduce to the end return x.permute(*range(len(noop)), *[len(noop) + i * 2 + 1 for i in range(len(i_))], *[len(noop) + i * 2 for i in range(len(i_))]) + + def unfold(self, dim:int, size, step:int) -> Self: + """ + Unfolds the tensor along dimension `dim` into overlapping windows. + + Each window has length `size` and begins every `step` elements of `self`. + Returns the input tensor with dimension `dim` replaced by dims `(n_windows, size)` + where `n_windows = (self.shape[dim] - size) // step + 1`. + + ```python exec="true" source="above" session="tensor" result="python" + unfolded = Tensor.arange(8).unfold(0,2,2) + print("\\n".join([repr(x.numpy()) for x in unfolded])) + ``` + ```python exec="true" source="above" session="tensor" result="python" + unfolded = Tensor.arange(27).reshape(3,3,3).unfold(-1,2,3) + print("\\n".join([repr(x.numpy()) for x in unfolded])) + ``` + """ + if size < 0: raise RuntimeError(f'size must be >= 0 but got {size=}') + if step <= 0: raise RuntimeError(f'step must be > 0 but got {step=}') + if size > self.shape[dim]: raise RuntimeError(f'maximum size for tensor at dimension {dim} is {self.shape[dim]} but size is {size}') + dim = self._resolve_dim(dim) + perm_to_last = tuple(i for i in range(self.ndim) if i != dim) + (dim,) + return self.permute(perm_to_last)._pool((size,), step).permute(argsort(perm_to_last) + (self.ndim,)) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 07ce5e217d..4f33052b65 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1447,30 +1447,6 @@ class Tensor(OpMixin): assert chunks > 0, f"expect chunks to be greater than 0, got: {chunks}" return list(self.split(ceildiv(dim_sz, chunks) if dim_sz else [0]*chunks, dim=dim)) - def unfold(self, dim:int, size:sint, step:int) -> Tensor: - """ - Unfolds the tensor along dimension `dim` into overlapping windows. - - Each window has length `size` and begins every `step` elements of `self`. - Returns the input tensor with dimension `dim` replaced by dims `(n_windows, size)` - where `n_windows = (self.shape[dim] - size) // step + 1`. - - ```python exec="true" source="above" session="tensor" result="python" - unfolded = Tensor.arange(8).unfold(0,2,2) - print("\\n".join([repr(x.numpy()) for x in unfolded])) - ``` - ```python exec="true" source="above" session="tensor" result="python" - unfolded = Tensor.arange(27).reshape(3,3,3).unfold(-1,2,3) - print("\\n".join([repr(x.numpy()) for x in unfolded])) - ``` - """ - if size < 0: raise RuntimeError(f'size must be >= 0 but got {size=}') - if step <= 0: raise RuntimeError(f'step must be > 0 but got {step=}') - if size > self.shape[dim]: raise RuntimeError(f'maximum size for tensor at dimension {dim} is {self.shape[dim]} but size is {size}') - dim = self._resolve_dim(dim) - perm_to_last = tuple(i for i in range(self.ndim) if i != dim) + (dim,) - return self.permute(perm_to_last)._pool((size,), step).permute(argsort(perm_to_last) + (self.ndim,)) - def meshgrid(self:Tensor, *args:Tensor, indexing:str="ij") -> tuple[Tensor, ...]: """ Generates coordinate matrices from coordinate vectors. @@ -2870,30 +2846,6 @@ class Tensor(OpMixin): """ return self._apply_uop(UOp.contiguous_backward) - def log(self) -> Tensor: - """ - Computes the natural logarithm element-wise. - - See: https://en.wikipedia.org/wiki/Logarithm - - ```python exec="true" source="above" session="tensor" result="python" - print(Tensor([1., 2., 4., 8.]).log().numpy()) - ``` - """ - return self.log2()*math.log(2) - - def log10(self) -> Tensor: - """ - Computes the base-10 logarithm element-wise. - - See: https://en.wikipedia.org/wiki/Logarithm - - ```python exec="true" source="above" session="tensor" result="python" - print(Tensor([1., 2., 4., 8.]).log10().numpy()) - ``` - """ - return self.log2()*math.log10(2) - def log2(self) -> Tensor: """ Computes the base-2 logarithm element-wise. @@ -3021,16 +2973,6 @@ class Tensor(OpMixin): # ***** math functions ***** - def round(self: Tensor) -> Tensor: - """ - Rounds the tensor element-wise with rounding half to even. - - ```python exec="true" source="above" session="tensor" result="python" - print(Tensor([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]).round().numpy()) - ``` - """ - return ((self > 0) == ((b := self.trunc() / 2.0).trunc() == b)).where((self - 0.5).ceil(), (self + 0.5).floor()) - def lerp(self, end:Tensor, weight:Tensor|float) -> Tensor: """ Linearly interpolates between `self` and `end` by `weight`. @@ -3136,42 +3078,6 @@ class Tensor(OpMixin): """ return (self.exp() + self.neg().exp()) / 2 - def atanh(self) -> Tensor: - """ - Applies the Inverse Hyperbolic Tangent (atanh) function element-wise. - - - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#atanh - - ```python exec="true" source="above" session="tensor" result="python" - print(Tensor([-0.9, -0.6, -0.3, 0., 0.3, 0.6, 0.9]).atanh().numpy()) - ``` - """ - return ((1 + self)/(1 - self)).log() / 2 - - def asinh(self) -> Tensor: - """ - Applies the Inverse Hyperbolic Sine (asinh) function element-wise. - - - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#asinh - - ```python exec="true" source="above" session="tensor" result="python" - print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).asinh().numpy()) - ``` - """ - return (self + (self.square() + 1).sqrt()).log() - - def acosh(self) -> Tensor: - """ - Applies the Inverse Hyperbolic Cosine (acosh) function element-wise. - - - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#acosh - - ```python exec="true" source="above" session="tensor" result="python" - print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).acosh().numpy()) - ``` - """ - return (self + (self.square() - 1).sqrt()).log() - def erf(self) -> Tensor: """ Applies error function element-wise.