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

python 3.10 is minimum [pr] (#7636)

Browse Source
This commit is contained in:
George Hotz
2024-11-11 23:05:50 +08:00
committed by GitHub
parent 6a0ed46b1c
commit aaa8059aec
2 changed files with 4 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tinygrad/helpers.py
View File

@@ -2,9 +2,7 @@ from __future__ import annotations
import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip
import urllib.request, subprocess, shutil, math, contextvars, types, copyreg, inspect, importlib
from dataclasses import dataclass
from typing import Dict, Tuple, Union, List, ClassVar, Optional, Iterable, Any, TypeVar, TYPE_CHECKING, Callable, Sequence
if TYPE_CHECKING: # TODO: remove this and import TypeGuard from typing once minimum python supported version is 3.10
from typing_extensions import TypeGuard
from typing import Dict, Tuple, Union, List, ClassVar, Optional, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard
T = TypeVar("T")
U = TypeVar("U")
@@ -320,10 +318,7 @@ class trange(tqdm):
def _reconstruct_code(*args): return types.CodeType(*args)
def _serialize_code(code:types.CodeType):
# NOTE: this works in Python 3.8 and up
if sys.version_info >= (3, 10): args = inspect.signature(types.CodeType).parameters
else: args = ['argcount', 'posonlyargcount', 'kwonlyargcount', 'nlocals', 'stacksize', 'flags', 'codestring',
'constants', 'names', 'varnames', 'filename', 'name', 'firstlineno', 'lnotab', 'freevars', 'cellvars']
args = inspect.signature(types.CodeType).parameters # NOTE: this works in Python 3.10 and up
return _reconstruct_code, tuple(code.__getattribute__('co_'+x.replace('codestring', 'code').replace('constants', 'consts')) for x in args)
copyreg.pickle(types.CodeType, _serialize_code)

5
tinygrad/ops.py
View File

@@ -384,7 +384,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
def const_factor(self) -> int:
"""largest known int that divides self"""
if self.op is Ops.CONST: return self.arg
if self.op is Ops.VCONST: return functools.reduce(math.gcd, self.arg)
if self.op is Ops.VCONST: return math.gcd(*self.arg)
if self.op is BinaryOps.ADD: return math.gcd(self.src[0].const_factor(), self.src[1].const_factor())
if self.op is BinaryOps.MUL: return self.src[0].arg if self.src[0].op is Ops.CONST else self.src[1].arg if self.src[1].op is Ops.CONST else 1
return 1
@@ -635,7 +635,6 @@ class PatternMatcher:
for p,fxn in self.patterns:
assert p.op is not None
tuple_fxn = fxn if isinstance(fxn, tuple) else deconstruct_function(fxn)
tuple_fxn[1]['__builtins__'] = __builtins__ # NOTE: Python 3.8 requires this for "all" and "len" and friends
real_fxn = types.FunctionType(*tuple_fxn)
for uop in p.op: self.pdict.setdefault(uop, []).append((p, real_fxn, p.early_reject, 'ctx' in inspect.signature(real_fxn).parameters))
@@ -906,7 +905,7 @@ def div_folding(x:UOp, c:int) -> Optional[UOp]:
def lt_folding(x:UOp, c:int) -> Optional[UOp]:
p, np = partition(split_uop(x, BinaryOps.ADD), lambda u: u.const_factor() == 1)
if np and (d:=functools.reduce(math.gcd, [u.const_factor() for u in np], c)) > 1 and 0 <= sum(u.vmin for u in p) and sum(u.vmax for u in p) < d:
if np and (d:=math.gcd(*[u.const_factor() for u in np], c)) > 1 and 0 <= sum(u.vmin for u in p) and sum(u.vmax for u in p) < d:
return cast(UOp, functools.reduce(operator.add, np).divides(d)).lt(c//d)
return None