github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-08 22:48:25 -05:00
Code Issues Packages Projects Releases Wiki Activity

Self type + mixins (#13056)

Browse Source 
* use Self type

* mixin

* fix later
This commit is contained in:
George Hotz
2025-11-02 13:30:01 +08:00
committed by GitHub
parent 8cbef912d2
commit 036ee9f84c
4 changed files with 66 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
.github/workflows/test.yml vendored
View File

@@ -243,8 +243,9 @@ jobs:
run: | run: |
python -m mypy --strict-equality --lineprecision-report . python -m mypy --strict-equality --lineprecision-report .
cat lineprecision.txt cat lineprecision.txt
- name: Run TYPED=1 # broken because of UPatAny
run: TYPED=1 python -c "import tinygrad" #- name: Run TYPED=1
# run: TYPED=1 python -c "import tinygrad"
unittest: unittest:
name: Unit Tests name: Unit Tests

4
tinygrad/tensor.py
View File

@@ -9,7 +9,7 @@ from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_u
from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, DEBUG, is_numpy_ndarray, FUSE_ATTENTION, SPEC from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, DEBUG, is_numpy_ndarray, FUSE_ATTENTION, SPEC
from tinygrad.helpers import suppress_finalizing from tinygrad.helpers import suppress_finalizing
from tinygrad.gradient import compute_gradient from tinygrad.gradient import compute_gradient
from tinygrad.uop.mathtraits import MathTrait from tinygrad.uop.mixins import MathMixin, MovementMixin
from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, identity_element, all_metadata, _index_to_concrete_int, sint_to_uop
from tinygrad.uop.spec import type_verify, tensor_spec from tinygrad.uop.spec import type_verify, tensor_spec
from tinygrad.device import Device, Buffer from tinygrad.device import Device, Buffer
@@ -100,7 +100,7 @@ def _flat_to_grouped(padding:Sequence[sint]) -> tuple[tuple[sint, sint], ...]: r
ReductionStr = Literal["mean", "sum", "none"] ReductionStr = Literal["mean", "sum", "none"]
class Tensor(MathTrait): class Tensor(MathMixin, MovementMixin):
""" """
A `Tensor` is a multi-dimensional matrix containing elements of a single data type. A `Tensor` is a multi-dimensional matrix containing elements of a single data type.

116
tinygrad/uop/mathtraits.py → tinygrad/uop/mixins.py
View File

@@ -1,4 +1,5 @@
from typing import TypeVar, TypeAlias, TYPE_CHECKING # mixins add syntactic sugar to Tensor and UOp
from typing import TypeAlias, TYPE_CHECKING, Self
from tinygrad.uop import Ops from tinygrad.uop import Ops
from tinygrad.dtype import dtypes, ConstType from tinygrad.dtype import dtypes, ConstType
from tinygrad.helpers import prod, argfix from tinygrad.helpers import prod, argfix
@@ -6,15 +7,14 @@ if TYPE_CHECKING:
from tinygrad.uop.ops import UOp from tinygrad.uop.ops import UOp
sint:TypeAlias = UOp|int sint:TypeAlias = UOp|int
TMT = TypeVar("TMT", bound="MathTrait") class MathMixin:
class MathTrait:
# required to implement # required to implement
def alu(self:TMT, op:Ops, *src:TMT) -> TMT: raise NotImplementedError def alu(self, op:Ops, *src:Self) -> Self: raise NotImplementedError
def const_like(self:TMT, b:ConstType) -> TMT: raise NotImplementedError def const_like(self, b:ConstType) -> Self: raise NotImplementedError
# great functions you get! # great functions you get!
def ufix(self:TMT, x:TMT|ConstType) -> TMT: return self.const_like(x) if not isinstance(x, MathTrait) else x def ufix(self, x:Self|ConstType) -> Self: return self.const_like(x) if not isinstance(x, MathMixin) else x
def _binop(self:TMT, op:Ops, x:TMT|ConstType, reverse:bool) -> TMT: def _binop(self, op:Ops, x:Self|ConstType, reverse:bool) -> Self:
return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x)) 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 logical_not(self): return self.ne(True)
def neg(self): def neg(self):
@@ -24,7 +24,7 @@ class MathTrait:
if (dtype:=getattr(self, 'dtype')) is not None: if (dtype:=getattr(self, 'dtype')) is not None:
if isinstance(dtype, tuple): dtype = dtype[0] 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") if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)): raise RuntimeError(f"{dtype} is not supported")
def add(self:TMT, x:TMT|ConstType, reverse:bool=False): def add(self, x:Self|ConstType, reverse:bool=False):
""" """
Adds `self` and `x`. Adds `self` and `x`.
Equivalent to `self + x`. Equivalent to `self + x`.
@@ -42,7 +42,7 @@ class MathTrait:
``` ```
""" """
return self._binop(Ops.ADD, x, reverse) return self._binop(Ops.ADD, x, reverse)
def mul(self:TMT, x:TMT|ConstType, reverse:bool=False): def mul(self, x:Self|ConstType, reverse:bool=False):
""" """
Multiplies `self` and `x`. Multiplies `self` and `x`.
Equivalent to `self * x`. Equivalent to `self * x`.
@@ -61,7 +61,7 @@ class MathTrait:
``` ```
""" """
return self._binop(Ops.MUL, x, reverse) return self._binop(Ops.MUL, x, reverse)
def bitwise_and(self:TMT, x:TMT|ConstType, reverse:bool=False): def bitwise_and(self, x:Self|ConstType, reverse:bool=False):
""" """
Computes the bitwise AND of `self` and `x`. Computes the bitwise AND of `self` and `x`.
Equivalent to `self & x`. Equivalent to `self & x`.
@@ -75,7 +75,7 @@ class MathTrait:
""" """
self._check_dtype() self._check_dtype()
return self._binop(Ops.AND, x, reverse) return self._binop(Ops.AND, x, reverse)
def bitwise_or(self:TMT, x:TMT|ConstType, reverse:bool=False): def bitwise_or(self, x:Self|ConstType, reverse:bool=False):
""" """
Computes the bitwise OR of `self` and `x`. Computes the bitwise OR of `self` and `x`.
Equivalent to `self | x`. Equivalent to `self | x`.
@@ -89,7 +89,7 @@ class MathTrait:
""" """
self._check_dtype() self._check_dtype()
return self._binop(Ops.OR, x, reverse) return self._binop(Ops.OR, x, reverse)
def bitwise_xor(self:TMT, x:TMT|ConstType, reverse:bool=False): def bitwise_xor(self, x:Self|ConstType, reverse:bool=False):
""" """
Computes bitwise xor of `self` and `x`. Computes bitwise xor of `self` and `x`.
Equivalent to `self ^ x`. Equivalent to `self ^ x`.
@@ -104,7 +104,7 @@ class MathTrait:
""" """
self._check_dtype() self._check_dtype()
return self._binop(Ops.XOR, x, reverse) return self._binop(Ops.XOR, x, reverse)
def idiv(self:TMT, x:TMT|ConstType, reverse:bool=False): def idiv(self, x:Self|ConstType, reverse:bool=False):
""" """
Divides `self` by `x`. Divides `self` by `x`.
Equivalent to `self // x`. Equivalent to `self // x`.
@@ -116,78 +116,78 @@ class MathTrait:
``` ```
""" """
return self._binop(Ops.IDIV, x, reverse) return self._binop(Ops.IDIV, x, reverse)
def mod(self:TMT, x:TMT|ConstType, reverse:bool=False): return self._binop(Ops.MOD, x, reverse) def mod(self, x:Self|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 sub(self, x:Self|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): def div(self, x:Self|ConstType, reverse:bool=False):
return (self.ufix(x)*self.alu(Ops.RECIPROCAL)) if reverse else (self*self.ufix(x).alu(Ops.RECIPROCAL)) return (self.ufix(x)*self.alu(Ops.RECIPROCAL)) if reverse else (self*self.ufix(x).alu(Ops.RECIPROCAL))
def __neg__(self): return self.neg() def __neg__(self): return self.neg()
def __add__(self:TMT, x:TMT|ConstType): return self.add(x) def __add__(self, x:Self|ConstType): return self.add(x)
def __sub__(self:TMT, x:TMT|ConstType): return self.sub(x) def __sub__(self, x:Self|ConstType): return self.sub(x)
def __mul__(self:TMT, x:TMT|ConstType): return self.mul(x) def __mul__(self, x:Self|ConstType): return self.mul(x)
def __truediv__(self:TMT, x:TMT|ConstType): return self.div(x) def __truediv__(self, x:Self|ConstType): return self.div(x)
def __floordiv__(self:TMT, x:TMT|ConstType): return self.idiv(x) # TODO: idiv is trunc div, not floordiv def __floordiv__(self, x:Self|ConstType): return self.idiv(x) # TODO: idiv is trunc div, not floordiv
def __mod__(self:TMT, x:TMT|ConstType): return self.mod(x) def __mod__(self, x:Self|ConstType): return self.mod(x)
def __and__(self:TMT, x:TMT|ConstType): return self.bitwise_and(x) def __and__(self, x:Self|ConstType): return self.bitwise_and(x)
def __or__(self:TMT, x:TMT|ConstType): return self.bitwise_or(x) def __or__(self, x:Self|ConstType): return self.bitwise_or(x)
def __xor__(self:TMT, x:TMT|ConstType): return self.bitwise_xor(x) def __xor__(self, x:Self|ConstType): return self.bitwise_xor(x)
def __radd__(self:TMT, x:TMT|ConstType): return self.add(x, True) def __radd__(self, x:Self|ConstType): return self.add(x, True)
def __rsub__(self:TMT, x:TMT|ConstType): return self.sub(x, True) def __rsub__(self, x:Self|ConstType): return self.sub(x, True)
def __rmul__(self:TMT, x:TMT|ConstType): return self.mul(x, True) def __rmul__(self, x:Self|ConstType): return self.mul(x, True)
def __rtruediv__(self:TMT, x:TMT|ConstType): return self.div(x, True) def __rtruediv__(self, x:Self|ConstType): return self.div(x, True)
def __rfloordiv__(self:TMT, x:TMT|ConstType): return self.idiv(x, True) def __rfloordiv__(self, x:Self|ConstType): return self.idiv(x, True)
def __rand__(self:TMT, x:TMT|ConstType): return self.bitwise_and(x, True) def __rand__(self, x:Self|ConstType): return self.bitwise_and(x, True)
def __ror__(self:TMT, x:TMT|ConstType): return self.bitwise_or(x, True) def __ror__(self, x:Self|ConstType): return self.bitwise_or(x, True)
def __rxor__(self:TMT, x:TMT|ConstType): return self.bitwise_xor(x, True) def __rxor__(self, x:Self|ConstType): return self.bitwise_xor(x, True)
def __rmod__(self:TMT, x:TMT|ConstType): return self.mod(x, True) def __rmod__(self, x:Self|ConstType): return self.mod(x, True)
def __lt__(self:TMT, x:TMT|ConstType): return self.alu(Ops.CMPLT, self.ufix(x)) def __lt__(self, x:Self|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 __gt__(self, x:Self|ConstType): return self.ufix(x).alu(Ops.CMPLT, self)
def __ge__(self:TMT, x:TMT|ConstType): return (self < x).logical_not() def __ge__(self, x:Self|ConstType): return (self < x).logical_not()
def __le__(self:TMT, x:TMT|ConstType): return (self > x).logical_not() def __le__(self, x:Self|ConstType): return (self > x).logical_not()
def ne(self:TMT, x:TMT|ConstType): return self.alu(Ops.CMPNE, self.ufix(x)) def ne(self, x:Self|ConstType): return self.alu(Ops.CMPNE, self.ufix(x))
def eq(self:TMT, x:TMT|ConstType): return self.ne(x).logical_not() def eq(self, x:Self|ConstType): return self.ne(x).logical_not()
def __ne__(self:TMT, x:TMT|ConstType): return self.ne(x) # type: ignore[override] def __ne__(self, x:Self|ConstType): return self.ne(x) # type: ignore[override]
# NOTE: __eq__ isn't overridden, and means the same thing as is by default # NOTE: __eq__ isn't overridden, and means the same thing as is by default
def lshift(self:TMT, x:TMT|int, reverse:bool=False): return self._binop(Ops.SHL, x, reverse) def lshift(self, x:Self|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 rshift(self, x:Self|int, reverse:bool=False): return self._binop(Ops.SHR, x, reverse)
def __lshift__(self:TMT, x:TMT|int): return self.lshift(x) def __lshift__(self, x:Self|int): return self.lshift(x)
def __rshift__(self:TMT, x:TMT|int): return self.rshift(x) def __rshift__(self, x:Self|int): return self.rshift(x)
def __rlshift__(self:TMT, x:TMT|int): return self.lshift(x, True) def __rlshift__(self, x:Self|int): return self.lshift(x, True)
def __rrshift__(self:TMT, x:TMT|int): return self.rshift(x, True) def __rrshift__(self, x:Self|int): return self.rshift(x, True)
def maximum(self:TMT, x:TMT|ConstType): return self.alu(Ops.MAX, self.ufix(x)) def maximum(self, x:Self|ConstType): return self.alu(Ops.MAX, self.ufix(x))
def minimum(self:TMT, x:TMT|ConstType): return -(-self).maximum(-x) def minimum(self, x:Self|ConstType): return -(-self).maximum(-x)
def where(self:TMT, x:TMT|ConstType, y:TMT|ConstType): def where(self, x:Self|ConstType, y:Self|ConstType):
if isinstance(x, type(self)): return self.alu(Ops.WHERE, x, x.ufix(y)) 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) if isinstance(y, type(self)): return self.alu(Ops.WHERE, y.ufix(x), y)
raise RuntimeError("where needs at least one UOp arg") raise RuntimeError("where needs at least one UOp arg")
def threefry(self:TMT, seed:TMT): return self.alu(Ops.THREEFRY, seed) def threefry(self, seed:Self): return self.alu(Ops.THREEFRY, seed)
def reciprocal(self): return self.alu(Ops.RECIPROCAL) def reciprocal(self): return self.alu(Ops.RECIPROCAL)
def trunc(self): return self.alu(Ops.TRUNC) def trunc(self): return self.alu(Ops.TRUNC)
def sqrt(self): return self.alu(Ops.SQRT) def sqrt(self): return self.alu(Ops.SQRT)
def sin(self): return self.alu(Ops.SIN) def sin(self): return self.alu(Ops.SIN)
def log2(self): return self.alu(Ops.LOG2) def log2(self): return self.alu(Ops.LOG2)
def exp2(self): return self.alu(Ops.EXP2) def exp2(self): return self.alu(Ops.EXP2)
def pow(self:TMT, x:TMT|ConstType): return self.alu(Ops.POW, self.ufix(x)) def pow(self, x:Self|ConstType): return self.alu(Ops.POW, self.ufix(x))
def __pow__(self:TMT, x:TMT|ConstType): return self.pow(x) def __pow__(self, x:Self|ConstType): return self.pow(x)
# **** movement ops ****
class MovementMixin:
# required to implement # required to implement
def _mop(self:TMT, op:Ops, arg) -> TMT: raise NotImplementedError def _mop(self, op:Ops, arg) -> Self: raise NotImplementedError
@property @property
def shape(self) -> tuple["sint", ...]: raise NotImplementedError def shape(self) -> tuple["sint", ...]: raise NotImplementedError
def view(self:TMT, shape, *args) -> TMT: # great functions you get!
def view(self, shape, *args) -> Self:
"""`.view` is an alias for `.reshape`.""" """`.view` is an alias for `.reshape`."""
return self.reshape(shape, *args) return self.reshape(shape, *args)
def reshape(self:TMT, shape, *args) -> TMT: def reshape(self, shape, *args) -> Self:
""" """
Returns a tensor with the same data as the original tensor but with a different shape. Returns a tensor with the same data as the original tensor but with a different shape.
`shape` can be passed as a tuple or as separate arguments. `shape` can be passed as a tuple or as separate arguments.

6
tinygrad/uop/ops.py
View File

@@ -4,7 +4,7 @@ import sys, time, functools, itertools, math, operator, hashlib, os, types, pick
from dataclasses import dataclass from dataclasses import dataclass
from enum import Enum, auto from enum import Enum, auto
from tinygrad.uop import Ops, GroupOp from tinygrad.uop import Ops, GroupOp
from tinygrad.uop.mathtraits import MathTrait from tinygrad.uop.mixins import MathMixin, MovementMixin
from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType, AddrSpace from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType, AddrSpace
from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC, CI from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC, CI
@@ -104,7 +104,7 @@ class recursive_property(property):
# NOTE: this should be frozen, but frozen is slower # NOTE: this should be frozen, but frozen is slower
@dataclass(eq=False, slots=True) @dataclass(eq=False, slots=True)
class UOp(MathTrait, metaclass=UOpMetaClass): class UOp(MathMixin, MovementMixin, metaclass=UOpMetaClass):
op:Ops op:Ops
dtype:DType = dtypes.void dtype:DType = dtypes.void
src:tuple[UOp, ...] = tuple() src:tuple[UOp, ...] = tuple()
@@ -853,7 +853,7 @@ def printable(loc:tuple[str, int]) -> str:
try: return lines(loc[0])[loc[1]-1].strip() try: return lines(loc[0])[loc[1]-1].strip()
except FileNotFoundError: return "<missing>" except FileNotFoundError: return "<missing>"
class UPat(MathTrait): class UPat(MathMixin, MovementMixin):
__slots__ = ("op", "dtype", "arg", "name", "src") __slots__ = ("op", "dtype", "arg", "name", "src")
def __init__(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|None=None, def __init__(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|None=None,
src:tuple[UPat, ...]|list[UPat]|UPat|None=None, arg:Any=None, src:tuple[UPat, ...]|list[UPat]|UPat|None=None, arg:Any=None,