""" btt_np: a pure NumPy Fourier KAN that runs anywhere, including in the browser (Pyodide). Three evaluation paths for the same trained weights: f32 float32 on the CPU bf16 inputs, weights and outputs rounded to bfloat16 (rounding only) toolkit a NumPy re-implementation of the Q.ANT toolkit CPU emulator's calc_kan_layer arithmetic: every multiply, add and cosine rounded to bfloat16 in the same order as the toolkit source (v2.3.1). Checked bit for bit against the real emulator. It is NOT the hardware: a card computes the optical "tcos" shape, not an exact cosine. Layer maths (one Fourier KAN layer): y_o = bias_o + sum_i sum_k a_oik * cos(k * x_i + phi_oik), k = 1..grid """ import math import numpy as np # ---------------------------------------------------------------- bfloat16 def bf16(x): """Round to the nearest bfloat16 (ties to even), returned as float32.""" shape = np.shape(x) a = np.ascontiguousarray(np.asarray(x, dtype=np.float32)).reshape(-1) u = a.view(np.uint32).astype(np.uint64) lsb = (u >> 16) & 1 r = ((u + 0x7FFF + lsb) >> 16) << 16 out = r.astype(np.uint32).view(np.float32) # keep NaN and infinity untouched special = ~np.isfinite(a) if special.any(): out = np.where(special, a, out) return out.reshape(shape) # ---------------------------------------------------------------- targets TARGETS = { "chirp": lambda x: np.sin(2 * np.pi * x) + np.cos(np.pi * x**2) + np.cos(np.pi * x**3) * np.sin(np.pi * x**3), "sine": lambda x: np.sin(2 * np.pi * x), "bumps": lambda x: np.exp(-4 * (x - 0.8) ** 2) - np.exp(-6 * (x + 1.0) ** 2) + 0.3 * np.cos(3 * x), } # ---------------------------------------------------------------- model class FourierKAN: """A stack of Fourier KAN layers with hand-written backprop and AdamW.""" def __init__(self, layer_sizes, grid_sizes, seed=0): if len(grid_sizes) != len(layer_sizes) - 1: raise ValueError("grid_sizes needs one entry per layer") rng = np.random.default_rng(seed) self.layers = [] for i in range(len(layer_sizes) - 1): n_in, n_out, g = layer_sizes[i], layer_sizes[i + 1], grid_sizes[i] std = math.sqrt(2.0 / ((n_in + n_out) * g)) self.layers.append({ "a": (rng.standard_normal((n_out, n_in, g)) * std).astype(np.float32), "p": rng.uniform(-math.pi / 2, math.pi / 2, (n_out, n_in, g)).astype(np.float32), "b": np.zeros(n_out, np.float32), "k": np.arange(1, g + 1, dtype=np.float32), }) self.layer_sizes, self.grid_sizes = list(layer_sizes), list(grid_sizes) @property def params(self): return sum(L["a"].size + L["p"].size + L["b"].size for L in self.layers) def forward(self, x, noise_std=0.0, rng=None, keep=False): """x: (batch, n_in). Returns output and, if keep, the cache for backprop.""" cache = [] h = x.astype(np.float32) for L in self.layers: theta = h[:, None, :, None] * L["k"] + L["p"][None] # (B, out, in, g) c = np.cos(theta) y = np.einsum("boig,oig->bo", c, L["a"]) + L["b"] if noise_std and rng is not None: y = y + rng.standard_normal(y.shape).astype(np.float32) * noise_std if keep: cache.append((h, theta, c)) h = y.astype(np.float32) return h, cache def backward(self, cache, grad_out): grads = [] g = grad_out.astype(np.float32) # (B, out) for L, (h, theta, c) in zip(reversed(self.layers), reversed(cache)): s = np.sin(theta) ga = np.einsum("bo,boig->oig", g, c) gp = -np.einsum("bo,boig->oig", g, s * L["a"][None]) gb = g.sum(0) gh = -np.einsum("bo,boig->bi", g, s * L["a"][None] * L["k"]) grads.append({"a": ga, "p": gp, "b": gb}) g = gh return list(reversed(grads)) # evaluation paths ----------------------------------------------------- def predict_f32(self, x): return self.forward(x)[0] def predict_bf16(self, x): h = bf16(x) for L in self.layers: a, p, k = bf16(L["a"]), bf16(L["p"]), bf16(L["k"]) y = np.einsum("boig,oig->bo", np.cos(h[:, None, :, None] * k + p[None]), a) h = bf16(bf16(y) + bf16(L["b"])) return h def predict_toolkit(self, x): h = bf16(x) for L in self.layers: y = toolkit_kan_layer(h, L["p"], L["a"], L["k"]) h = bf16(y + bf16(L["b"])) # bias added on the host, as in the official examples return h def toolkit_kan_layer(x, phis, ampls, ks): """Same arithmetic, same rounding order as the toolkit's calc_kan_layer on its CPU backend.""" x, phis, ampls, ks = bf16(x), bf16(phis), bf16(ampls), bf16(ks) n_out, n_in, n_k = phis.shape kx = bf16(x[:, :, None] * ks[None, None, :]) # (B, in, k) multiply on device acc = np.zeros((x.shape[0], n_out), np.float32) for i in range(n_in): # sum order: in, then k for k in range(n_k): arg = bf16(kx[:, i, k][:, None] + phis[None, :, i, k]) # add on host, bf16 term = bf16(bf16(np.cos(arg)) * ampls[None, :, i, k]) # tcos(u) * v on device acc = bf16(acc + term) return acc # ---------------------------------------------------------------- training def train(cfg, on_progress=None, every=None): """cfg keys: target, layer_sizes, grid_sizes, lr, epochs, batch, noise_std, seed, scheduler.""" f = TARGETS[cfg["target"]] model = FourierKAN(cfg["layer_sizes"], cfg["grid_sizes"], cfg.get("seed", 0)) rng = np.random.default_rng(42 + cfg.get("seed", 0)) lr0, steps, batch = cfg["lr"], cfg["epochs"], cfg.get("batch", 128) noise = cfg.get("noise_std", 0.0) or 0.0 b1, b2, eps, wd = 0.9, 0.999, 1e-8, 0.01 # PyTorch AdamW defaults m = [{k: np.zeros_like(L[k]) for k in "apb"} for L in model.layers] v = [{k: np.zeros_like(L[k]) for k in "apb"} for L in model.layers] losses = [] every = every or max(1, steps // 50) for t in range(1, steps + 1): lr = lr0 * 0.5 * (1 + math.cos(math.pi * (t - 1) / steps)) if cfg.get("scheduler") == "cosine" else lr0 xb = rng.uniform(-2, 2, (batch, 1)).astype(np.float32) yb = f(xb).astype(np.float32) out, cache = model.forward(xb, noise, rng, keep=True) diff = out - yb loss = float(np.mean(diff * diff)) losses.append(loss) grads = model.backward(cache, 2.0 * diff / diff.size) for li, (L, G) in enumerate(zip(model.layers, grads)): for k in "apb": L[k] *= (1 - lr * wd) m[li][k] = b1 * m[li][k] + (1 - b1) * G[k] v[li][k] = b2 * v[li][k] + (1 - b2) * G[k] * G[k] mh = m[li][k] / (1 - b1**t) vh = v[li][k] / (1 - b2**t) L[k] -= (lr * mh / (np.sqrt(vh) + eps)).astype(np.float32) if on_progress and (t % every == 0 or t == steps): on_progress(t, float(np.mean(losses[-every:]))) return model, losses def evaluate(model, target="chirp", n=500): f = TARGETS[target] x = np.linspace(-2, 2, n, dtype=np.float32) y = f(x).astype(np.float32) X = x[:, None] out = {} for name, fn in (("f32", model.predict_f32), ("bf16", model.predict_bf16), ("toolkit", model.predict_toolkit)): p = fn(X).squeeze() out[name] = {"mse": float(np.mean((y - p) ** 2)), "pred": p} return x, y, out, float(np.mean(y.astype(np.float64) ** 2))