File size: 2,235 Bytes
ae01165 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | #!/usr/bin/env python3
"""GaugeBench. Exact Q rate on listed polys. PROCESS vs float AD vs finite diff."""
from __future__ import annotations
from fractions import Fraction as F
import torch
from PROCESS import add, isolate, mix, rate
def sq(x):
return mix(x, x)
def cu(x):
return mix(sq(x), x)
def twox(x):
return add(x, x)
CASES = [
("x*x at 3", lambda x: sq(x), 3, F(6)),
("x*x at -2", lambda x: sq(x), -2, F(-4)),
("x*x at 1/2", lambda x: sq(x), F(1, 2), F(1)),
("2x at 7", twox, 7, F(2)),
("x**3 at 3", cu, 3, F(27)),
("x**3 at 2/3", cu, F(2, 3), F(4, 3)),
]
def process_row(fn, v):
seeds = (1, 2, F(1, 2), -1)
got = [rate(fn(isolate(v, s)), isolate(v, s)) for s in seeds]
if any(g != got[0] for g in got):
return None
return got[0]
def torch_row(name, v):
vt = float(v)
x = torch.tensor(vt, requires_grad=True)
if name.startswith("x*x"):
y = x * x
elif name.startswith("2x"):
y = x + x
elif name.startswith("x**3"):
y = x * x * x
else:
return None
y.backward()
return float(x.grad)
def fd_row(name, v, h=1e-6):
vt = float(v)
def f(z):
if name.startswith("x*x"):
return z * z
if name.startswith("2x"):
return z + z
if name.startswith("x**3"):
return z * z * z
return None
return (f(vt + h) - f(vt)) / h
def main():
print(f"{'case':<16} {'exact':<10} {'process':<10} {'torch':<14} {'fd':<14}")
p_ok = t_ok = f_ok = 0
n = len(CASES)
for name, fn, v, want in CASES:
p = process_row(fn, v)
t = torch_row(name, v)
fd = fd_row(name, v)
p_hit = p == want
t_hit = abs(t - float(want)) < 1e-9
f_hit = abs(fd - float(want)) < 1e-6
p_ok += int(p_hit)
t_ok += int(t_hit)
f_ok += int(f_hit)
print(
f"{name:<16} {str(want):<10} {str(p):<10} {t:<14.8g} {fd:<14.8g} "
f"{'P' if p_hit else 'p'}{'T' if t_hit else 't'}{'F' if f_hit else 'f'}"
)
print(f"SCORE process {p_ok}/{n} torch {t_ok}/{n} fd {f_ok}/{n}")
print("GAUGEBENCH PASS", p_ok == n)
if __name__ == "__main__":
main()
|