File size: 5,484 Bytes
d7070c6 | 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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | # SPDX-License-Identifier: CC-BY-NC-4.0
# Copyright (c) 2026 Jakub Ćwirlej
"""Unit tests for tethergate.py (CPU, a few seconds): python test_tethergate.py"""
import math
import sys
import torch
from tethergate import TetherGate, floor_schedule, tethergate_step
def logit(p):
return math.log(p / (1.0 - p))
def test_floor_schedule():
kw = dict(ramp_start=100, f_min=0.05, tau=50.0, ramp_end=200)
ok = floor_schedule(0, **kw) == 1.0 and floor_schedule(99, **kw) == 1.0
exp_vals = [floor_schedule(t, mode='exp', **kw) for t in range(100, 2000, 7)]
cos_vals = [floor_schedule(t, mode='cos', **kw) for t in range(100, 400, 3)]
ok &= all(a >= b for a, b in zip(exp_vals, exp_vals[1:])) and min(exp_vals) > 0.05
ok &= all(a >= b for a, b in zip(cos_vals, cos_vals[1:])) and abs(floor_schedule(200, mode='cos', **kw) - 0.05) < 1e-12
return ok
def test_forward_floor():
g = TetherGate(beta=6.0, theta_init=0.3)
s = torch.linspace(-2, 2, 11)
ok = torch.equal(g(s), torch.ones(11)) # floor 1.0 -> gate fully open
g.floor.fill_(0.0)
ref = torch.sigmoid((6.0 * (s - 0.3)).clamp(-8.0, 8.0))
ok &= torch.allclose(g(s), ref)
g.floor.fill_(0.2)
ok &= torch.allclose(g(s), 0.2 + 0.8 * ref)
return bool(ok)
def test_gradient_alive_below_floor():
"""g_raw below the floor but the argument inside (-8, 8): theta still receives a gradient."""
g = TetherGate(beta=6.0, theta_init=0.8) # arg = 6*(0-0.8) = -4.8 -> g_raw = 0.008 < floor
g.floor.fill_(0.05)
g(torch.tensor([0.0])).sum().backward()
return g.theta.grad is not None and float(g.theta.grad.abs()) > 0
def test_saturation_beyond_clamp_and_spring_recovery():
"""PITFALL (documented): with beta*(s - theta) < arg_min the clamp saturates and the gradient to
theta is exactly zero. The spring is not gradient-based, so it still pulls theta back."""
g = TetherGate(beta=6.0, g_target=0.5, theta_init=5.0) # arg = -30 -> clamped at -8
g.floor.fill_(0.05)
g(torch.tensor([0.0])).sum().backward()
dead = g.theta.grad is not None and float(g.theta.grad.abs()) == 0.0
s = torch.tensor([0.0])
for step in range(3000):
g(s)
tethergate_step([g], step, lr=1e-2, wd=0.5, ramp_start=0, f_min=0.05, mode='cos', ramp_end=1)
return dead and abs(g.theta.item() - g.theta_star()) < 1e-3
def test_detach_signal():
s = torch.tensor([0.5, 1.0], requires_grad=True)
TetherGate(detach_signal=True)(s).sum().backward()
a = s.grad is None
s2 = torch.tensor([0.5, 1.0], requires_grad=True)
g = TetherGate(detach_signal=False); g.floor.fill_(0.0)
g(s2).sum().backward()
return a and s2.grad is not None and float(s2.grad.abs().sum()) > 0
def test_ema_train_eval():
g = TetherGate(ema_momentum=0.9)
g.train(); g(torch.full((4,), 2.0))
a = abs(float(g.ema_s) - 2.0) < 1e-6 # first call initializes from NaN
g(torch.full((4,), 0.0))
b = abs(float(g.ema_s) - 1.8) < 1e-6
g.eval(); g(torch.full((4,), 100.0))
return a and b and abs(float(g.ema_s) - 1.8) < 1e-6 # eval does not update
def test_spring_reaches_target():
"""No loss pressure, constant signal: theta -> theta_star and the gate -> g_target."""
g = TetherGate(beta=6.0, g_target=0.35, theta_init=2.0)
s = torch.full((8,), 0.7)
for step in range(3000):
g(s)
tethergate_step([g], step, lr=1e-2, wd=0.5, ramp_start=0, f_min=0.0, mode='cos', ramp_end=1)
g.eval()
return abs(g.theta.item() - g.theta_star()) < 1e-3 and abs(g(s).mean().item() - 0.35) < 1e-3
def test_equilibrium_independent_of_lr():
"""Constant loss pressure r on theta: theta_eq - theta_star = r / wd, for any LR."""
r, wd, out = 0.02, 0.5, []
for lr in (1e-3, 1e-2):
g = TetherGate(beta=6.0, g_target=0.5, theta_init=0.0)
s = torch.full((4,), 0.4)
n = int(30 / (lr * wd)) # ~30 time constants
for step in range(n):
g(s)
with torch.no_grad():
g.theta.add_(lr * r) # SGD step with dL/dtheta = -r
tethergate_step([g], step, lr=lr, wd=wd, ramp_start=0, f_min=0.0, mode='cos', ramp_end=1)
out.append(g.theta.item() - g.theta_star())
return all(abs(d - r / wd) < 2e-3 for d in out)
def test_global_vs_local_anchor():
def run(anchor):
gs = [TetherGate(beta=6.0, g_target=0.5) for _ in range(2)]
for step in range(2000):
gs[0](torch.full((4,), 0.2)); gs[1](torch.full((4,), 1.0))
tethergate_step(gs, step, lr=1e-2, wd=0.5, anchor=anchor, ramp_start=0, f_min=0.0, mode='cos', ramp_end=1)
return [g.theta.item() for g in gs]
glob, loc = run('global'), run('local')
return abs(glob[0] - glob[1]) < 1e-3 and abs(glob[0] - 0.6) < 1e-3 and abs(loc[0] - 0.2) < 1e-3 and abs(loc[1] - 1.0) < 1e-3
if __name__ == '__main__':
tests = [test_floor_schedule, test_forward_floor, test_gradient_alive_below_floor,
test_saturation_beyond_clamp_and_spring_recovery, test_detach_signal,
test_ema_train_eval, test_spring_reaches_target, test_equilibrium_independent_of_lr,
test_global_vs_local_anchor]
res = [(t.__name__, bool(t())) for t in tests]
for n, r in res:
print(f' {n:38s} {"OK" if r else "FAIL"}')
sys.exit(0 if all(r for _, r in res) else 1)
|