From 1ecb99480ef03271c3fc55a63d102c1d0a47ae6a Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 14 Oct 2025 10:59:00 +0800 Subject: [PATCH] add typing to MathTraits --- tinygrad/apps/llm.py | 5 +- tinygrad/helpers.py | 2 +- tinygrad/nn/__init__.py | 8 +-- tinygrad/tensor.py | 5 +- tinygrad/uop/mathtraits.py | 105 +++++++++++++++++++------------------ tinygrad/uop/ops.py | 8 +-- 6 files changed, 68 insertions(+), 65 deletions(-) diff --git a/tinygrad/apps/llm.py b/tinygrad/apps/llm.py index a718170259..77fa753ec0 100644 --- a/tinygrad/apps/llm.py +++ b/tinygrad/apps/llm.py @@ -55,9 +55,10 @@ class SimpleTokenizer: def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor: B, H, T, Hd = x.shape - assert (Hd & 1) == 0, "RoPE requires an even head dimension" + assert isinstance(Hd, int) and (Hd & 1) == 0, "RoPE requires an even head dimension" half = Hd // 2 - angles = (Tensor.arange(T, dtype="float32") + start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :] + t_start_pos = start_pos if isinstance(start_pos, int) else Tensor(start_pos) + angles = (Tensor.arange(T, dtype="float32") + t_start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :] # contiguous here allows RoPE to be pruned in the JIT cos, sin = angles.cos().reshape(1, 1, T, half).cast(x.dtype).contiguous(), angles.sin().reshape(1, 1, T, half).cast(x.dtype).contiguous() x_pairs = x.reshape(B, H, T, half, 2) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 00e2de83d2..e70e927801 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -153,7 +153,7 @@ CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), Co ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0) FUSE_ATTENTION = ContextVar("FUSE_ATTENTION", 0) EMULATE = ContextVar("EMULATE", "") -CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if (aff:=getattr(os, "sched_getaffinity", None)) else (os.cpu_count() or 1))) +CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(aff(0)) if (aff:=getattr(os, "sched_getaffinity", None)) else (os.cpu_count() or 1))) CPU_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1) VIZ = PROFILE = ContextVar("VIZ", 0) SPEC = ContextVar("SPEC", 0) diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index d32a3d5e2f..b27ab036c0 100644 --- a/tinygrad/nn/__init__.py +++ b/tinygrad/nn/__init__.py @@ -223,7 +223,7 @@ class InstanceNorm: print(t.mean().item(), t.std().item()) ``` """ - def __init__(self, num_features:int, eps=1e-5, affine=True): + def __init__(self, num_features:int, eps:float=1e-5, affine:bool=True): self.num_features, self.eps = num_features, eps self.weight: Tensor|None = Tensor.ones(num_features) if affine else None self.bias: Tensor|None = Tensor.zeros(num_features) if affine else None @@ -249,16 +249,16 @@ class LayerNorm: print(t.mean().item(), t.std().item()) ``` """ - def __init__(self, normalized_shape:int|tuple[int, ...], eps=1e-5, elementwise_affine=True): + def __init__(self, normalized_shape:int|tuple[int, ...], eps:float=1e-5, elementwise_affine:bool=True): self.normalized_shape: tuple[int, ...] = make_tuple(normalized_shape, 1) - self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine + self.axis, self.eps = tuple(-1-i for i in range(len(self.normalized_shape))), eps self.weight: Tensor|None = Tensor.ones(*self.normalized_shape) if elementwise_affine else None self.bias: Tensor|None = Tensor.zeros(*self.normalized_shape) if elementwise_affine else None def __call__(self, x:Tensor) -> Tensor: assert self.normalized_shape == x.shape[-len(self.normalized_shape):], f"last dimensions of {x.shape} must match {self.normalized_shape}" x = x.layernorm(eps=self.eps, axis=self.axis) - if not self.elementwise_affine: return x + if self.weight is None or self.bias is None: return x return x * self.weight + self.bias class LayerNorm2d(LayerNorm): diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index fc40dfaac7..d11ccac0a7 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1212,6 +1212,7 @@ class Tensor(MathTrait): match index: case Tensor(): if not dtypes.is_int(index.dtype): raise IndexError(f"index dtype {index.dtype} is not supported") + assert isinstance(size, int), "size must be an int" index = (index < 0).where(index+size, index).to(self.device) # treat negative index values case list() | tuple(): if not dtypes.is_int((ti:=Tensor(index)).dtype): raise IndexError(f"{index=} contains non-int element") @@ -3723,7 +3724,7 @@ class Tensor(MathTrait): if self.dtype != dtypes.bool and not dtypes.is_int(self.dtype): raise RuntimeError(f"{self.dtype} is not supported") return self.logical_not() if self.dtype == dtypes.bool else self ^ -1 - def lshift(self, x:int, reverse=False) -> Tensor: + def lshift(self, x:Tensor|int, reverse=False) -> Tensor: """ Computes left arithmetic shift of `self` by `x` bits. `self` must have unsigned dtype. Equivalent to `self << x`. @@ -3735,7 +3736,7 @@ class Tensor(MathTrait): assert dtypes.is_unsigned(self.dtype) and isinstance(x, int) and x >= 0 and not reverse, f"not supported {self.dtype=} {x=}" return self.mul(2 ** x, reverse) - def rshift(self, x:int, reverse=False) -> Tensor: + def rshift(self, x:Tensor|int, reverse=False) -> Tensor: """ Computes right arithmetic shift of `self` by `x` bits. `self` must have unsigned dtype. Equivalent to `self >> x`. diff --git a/tinygrad/uop/mathtraits.py b/tinygrad/uop/mathtraits.py index 2da0ea887a..5897f79ee7 100644 --- a/tinygrad/uop/mathtraits.py +++ b/tinygrad/uop/mathtraits.py @@ -2,15 +2,15 @@ from typing import TypeVar from tinygrad.uop import Ops from tinygrad.dtype import dtypes, ConstType -TMathTrait = TypeVar("TMathTrait", bound="MathTrait") +TMT = TypeVar("TMT", bound="MathTrait") class MathTrait: # required to implement - def alu(self:TMathTrait, op:Ops, *src:TMathTrait) -> TMathTrait: raise NotImplementedError - def const_like(self:TMathTrait, b:ConstType) -> TMathTrait: raise NotImplementedError + def alu(self:TMT, op:Ops, *src:TMT) -> TMT: raise NotImplementedError + def const_like(self:TMT, b:ConstType) -> TMT: raise NotImplementedError # great functions you get! - def ufix(self:TMathTrait, x:ConstType|TMathTrait) -> TMathTrait: return self.const_like(x) if not isinstance(x, MathTrait) else x - def _binop(self:TMathTrait, op:Ops, x:TMathTrait|ConstType, reverse:bool) -> TMathTrait: + def ufix(self:TMT, x:TMT|ConstType) -> TMT: return self.const_like(x) if not isinstance(x, MathTrait) else x + def _binop(self:TMT, op:Ops, x:TMT|ConstType, reverse:bool) -> TMT: return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x)) def logical_not(self): return self.ne(True) def neg(self): @@ -20,7 +20,7 @@ class MathTrait: if (dtype:=getattr(self, 'dtype')) is not None: if isinstance(dtype, tuple): dtype = dtype[0] if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)): raise RuntimeError(f"{dtype} is not supported") - def add(self, x, reverse=False): + def add(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Adds `self` and `x`. Equivalent to `self + x`. @@ -38,7 +38,7 @@ class MathTrait: ``` """ return self._binop(Ops.ADD, x, reverse) - def mul(self, x, reverse=False): + def mul(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Multiplies `self` and `x`. Equivalent to `self * x`. @@ -57,7 +57,7 @@ class MathTrait: ``` """ return self._binop(Ops.MUL, x, reverse) - def bitwise_and(self, x, reverse=False): + def bitwise_and(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Computes the bitwise AND of `self` and `x`. Equivalent to `self & x`. @@ -71,7 +71,7 @@ class MathTrait: """ self._check_dtype() return self._binop(Ops.AND, x, reverse) - def bitwise_or(self, x, reverse=False): + def bitwise_or(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Computes the bitwise OR of `self` and `x`. Equivalent to `self | x`. @@ -85,7 +85,7 @@ class MathTrait: """ self._check_dtype() return self._binop(Ops.OR, x, reverse) - def bitwise_xor(self, x, reverse=False): + def bitwise_xor(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Computes bitwise xor of `self` and `x`. Equivalent to `self ^ x`. @@ -100,7 +100,7 @@ class MathTrait: """ self._check_dtype() return self._binop(Ops.XOR, x, reverse) - def idiv(self, x, reverse=False): + def idiv(self:TMT, x:TMT|ConstType, reverse:bool=False): """ Divides `self` by `x`. Equivalent to `self // x`. @@ -112,61 +112,62 @@ class MathTrait: ``` """ return self._binop(Ops.IDIV, x, reverse) - def mod(self, x, reverse=False): return self._binop(Ops.MOD, x, reverse) - def sub(self, x, reverse=False): return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x)) - def div(self, x, reverse=False): return (self.ufix(x)*self.alu(Ops.RECIP)) if reverse else (self*self.ufix(x).alu(Ops.RECIP)) + def mod(self:TMT, x:TMT|ConstType, reverse:bool=False): return self._binop(Ops.MOD, x, reverse) + def sub(self:TMT, x:TMT|ConstType, reverse:bool=False): return self.ufix(x).alu(Ops.ADD, -self) if reverse else self.alu(Ops.ADD, self.ufix(-x)) + def div(self:TMT, x:TMT|ConstType, reverse:bool=False): return (self.ufix(x)*self.alu(Ops.RECIP)) if reverse else (self*self.ufix(x).alu(Ops.RECIP)) def __neg__(self): return self.neg() - def __add__(self, x): return self.add(x) - def __sub__(self, x): return self.sub(x) - def __mul__(self, x): return self.mul(x) - def __truediv__(self, x): return self.div(x) - def __floordiv__(self, x): return self.idiv(x) # TODO: idiv is trunc div, not floordiv - def __mod__(self, x): return self.mod(x) - def __and__(self, x): return self.bitwise_and(x) - def __or__(self, x): return self.bitwise_or(x) - def __xor__(self, x): return self.bitwise_xor(x) + def __add__(self:TMT, x:TMT|ConstType): return self.add(x) + def __sub__(self:TMT, x:TMT|ConstType): return self.sub(x) + def __mul__(self:TMT, x:TMT|ConstType): return self.mul(x) + def __truediv__(self:TMT, x:TMT|ConstType): return self.div(x) + def __floordiv__(self:TMT, x:TMT|ConstType): return self.idiv(x) # TODO: idiv is trunc div, not floordiv + def __mod__(self:TMT, x:TMT|ConstType): return self.mod(x) + def __and__(self:TMT, x:TMT|ConstType): return self.bitwise_and(x) + def __or__(self:TMT, x:TMT|ConstType): return self.bitwise_or(x) + def __xor__(self:TMT, x:TMT|ConstType): return self.bitwise_xor(x) - def __radd__(self, x): return self.add(x, True) - def __rsub__(self, x): return self.sub(x, True) - def __rmul__(self, x): return self.mul(x, True) - def __rtruediv__(self, x): return self.div(x, True) - def __rfloordiv__(self, x): return self.idiv(x, True) - def __rand__(self, x): return self.bitwise_and(x, True) - def __ror__(self, x): return self.bitwise_or(x, True) - def __rxor__(self, x): return self.bitwise_xor(x, True) - def __rmod__(self, x): return self.mod(x, True) + def __radd__(self:TMT, x:TMT|ConstType): return self.add(x, True) + def __rsub__(self:TMT, x:TMT|ConstType): return self.sub(x, True) + def __rmul__(self:TMT, x:TMT|ConstType): return self.mul(x, True) + def __rtruediv__(self:TMT, x:TMT|ConstType): return self.div(x, True) + def __rfloordiv__(self:TMT, x:TMT|ConstType): return self.idiv(x, True) + def __rand__(self:TMT, x:TMT|ConstType): return self.bitwise_and(x, True) + def __ror__(self:TMT, x:TMT|ConstType): return self.bitwise_or(x, True) + def __rxor__(self:TMT, x:TMT|ConstType): return self.bitwise_xor(x, True) + def __rmod__(self:TMT, x:TMT|ConstType): return self.mod(x, True) - def __lt__(self, x): return self.alu(Ops.CMPLT, self.ufix(x)) - def __gt__(self, x): return self.ufix(x).alu(Ops.CMPLT, self) - def __ge__(self, x): return (self < x).logical_not() - def __le__(self, x): return (self > x).logical_not() + def __lt__(self:TMT, x:TMT|ConstType): return self.alu(Ops.CMPLT, self.ufix(x)) + def __gt__(self:TMT, x:TMT|ConstType): return self.ufix(x).alu(Ops.CMPLT, self) + def __ge__(self:TMT, x:TMT|ConstType): return (self < x).logical_not() + def __le__(self:TMT, x:TMT|ConstType): return (self > x).logical_not() - def ne(self, x): return self.alu(Ops.CMPNE, self.ufix(x)) - def eq(self, x): return self.ne(x).logical_not() + def ne(self:TMT, x:TMT|ConstType): return self.alu(Ops.CMPNE, self.ufix(x)) + def eq(self:TMT, x:TMT|ConstType): return self.ne(x).logical_not() + # TODO: make typing of __ne__ work def __ne__(self, x): return self.ne(x) # NOTE: __eq__ isn't overridden, and means the same thing as is by default - def lshift(self, x, reverse=False): return self._binop(Ops.SHL, x, reverse) - def rshift(self, x, reverse=False): return self._binop(Ops.SHR, x, reverse) - def __lshift__(self, x): return self.lshift(x) - def __rshift__(self, x): return self.rshift(x) - def __rlshift__(self, x): return self.lshift(x, True) - def __rrshift__(self, x): return self.rshift(x, True) + def lshift(self:TMT, x:TMT|int, reverse:bool=False): return self._binop(Ops.SHL, x, reverse) + def rshift(self:TMT, x:TMT|int, reverse:bool=False): return self._binop(Ops.SHR, x, reverse) + def __lshift__(self:TMT, x:TMT|int): return self.lshift(x) + def __rshift__(self:TMT, x:TMT|int): return self.rshift(x) + def __rlshift__(self:TMT, x:TMT|int): return self.lshift(x, True) + def __rrshift__(self:TMT, x:TMT|int): return self.rshift(x, True) - def maximum(self, x): return self.alu(Ops.MAX, self.ufix(x)) - def minimum(self, x): return -(-self).maximum(-x) - def where(self, x, y): - if type(self) is type(x): return self.alu(Ops.WHERE, x, x.ufix(y)) - if type(self) is type(y): return self.alu(Ops.WHERE, y.ufix(x), y) + def maximum(self:TMT, x:TMT|ConstType): return self.alu(Ops.MAX, self.ufix(x)) + def minimum(self:TMT, x:TMT|ConstType): return -(-self).maximum(-x) + def where(self:TMT, x:TMT|ConstType, y:TMT|ConstType): + if isinstance(x, type(self)): return self.alu(Ops.WHERE, x, x.ufix(y)) + if isinstance(y, type(self)): return self.alu(Ops.WHERE, y.ufix(x), y) raise RuntimeError("where needs at least one UOp arg") - def threefry(self, seed): return self.alu(Ops.THREEFRY, seed) + def threefry(self:TMT, seed:TMT): return self.alu(Ops.THREEFRY, seed) def reciprocal(self): return self.alu(Ops.RECIP) def trunc(self): return self.alu(Ops.TRUNC) def sqrt(self): return self.alu(Ops.SQRT) def sin(self): return self.alu(Ops.SIN) def log2(self): return self.alu(Ops.LOG2) def exp2(self): return self.alu(Ops.EXP2) - def pow(self, x): return self.alu(Ops.POW, self.ufix(x)) - def __pow__(self, x): return self.pow(x) + def pow(self:TMT, x:TMT|ConstType): return self.alu(Ops.POW, self.ufix(x)) + def __pow__(self:TMT, x:TMT|ConstType): return self.pow(x) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 16900e97fc..2f1b1d211e 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -34,11 +34,11 @@ def resolve(x:UOp|bool, default:bool=True): def _suop(lst, uop_fxn, python_fxn): uops, nums = partition(lst, lambda x: isinstance(x, UOp)) return ssimplify(functools.reduce(uop_fxn, uops + ([python_fxn(nums)] if nums else []))) -def smax(*lst): return _suop(argfix(*lst), UOp.maximum, max) -def smin(*lst): return _suop(argfix(*lst), UOp.minimum, min) -def srender(x) -> str: return x.render() if isinstance(x, UOp) else str(x) +def smax(*lst:sint): return _suop(argfix(*lst), UOp.maximum, max) +def smin(*lst:sint): return _suop(argfix(*lst), UOp.minimum, min) +def srender(x:sint) -> str: return x.render() if isinstance(x, UOp) else str(x) -def ssimplify(uop): return uop.ssimplify() if isinstance(uop, UOp) else uop +def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop def range_str(u:UOp) -> str: return '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]])