The walkthrough · Beyond the Transistor

The layer

A Fourier KAN layer holds two learnable tensors and one fixed one. Click any highlighted line for what it does and why.

btt/model.py

1class FourierKAN(torch.nn.Module):2    def __init__(self, in_features, out_features, grid, bias=True, noise_std=None):3        super().__init__()4        std = math.sqrt(2.0 / ((in_features + out_features) * grid))5        self.amplitude = torch.nn.Parameter(torch.randn(out_features, in_features, grid) * std)6        self.phase = torch.nn.Parameter(7            torch.empty(out_features, in_features, grid).uniform_(-math.pi / 2, math.pi / 2)8        )9        self.bias = torch.nn.Parameter(torch.zeros(out_features)) if bias else None10        self.register_buffer("k", torch.arange(1, grid + 1, dtype=torch.float32))11        self.noise_std = noise_std12 13    def forward(self, x):14        angles = x[:, None, :, None] * self.k + self.phase[None]15        y = (self.amplitude[None] * torch.cos(angles)).sum(dim=(2, 3))16        if self.training and self.noise_std:17            y = y + torch.randn_like(y) * self.noise_std18        if self.bias is not None:19            y = y + self.bias20        return y

Scale the starting amplitudes

Each output sums in × grid cosine terms. Without scaling, the variance of that sum grows with layer size and deep stacks blow up at step one. This is the same idea as Glorot initialisation, spread across the grid as well as the fan in and fan out.

Click a highlighted line.

Training and evaluation

The training loop is ordinary PyTorch. The interesting part is after it: the same weights are evaluated on several backends so the cost of each number format is visible.

btt/runner.py, condensed

1model = build_model(cfg["layer_sizes"], cfg["grid_sizes"], cfg["noise_std"])2opt = torch.optim.AdamW(model.parameters(), lr=cfg["lr"])3sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, cfg["epochs"])4for step in range(cfg["epochs"]):5    xb = rng.uniform(-2, 2, size=(cfg["batch"], 1))6    yb = torch.tensor(f(xb), dtype=torch.float32)7    opt.zero_grad()8    loss = torch.mean((model(torch.tensor(xb, dtype=torch.float32)) - yb) ** 2)9    loss.backward()10    opt.step()11    sched.step()12 13x = np.linspace(-2, 2, 500)14for name in cfg["backends"]:15    pred, why = execute(model, x, name)16    results[name] = np.mean((f(x) - pred) ** 2)

Stack layers into a network

layer_sizes like [1, 16, 1] means one input, one hidden layer of 16, one output. grid_sizes gives one grid per layer. The playground builds these two lists for you and counts the parameters live.

Click a highlighted line.

Run it yourself

The playground exports this exact code as a zip with a notebook inside. It runs on any CPU in under a minute for small networks, and on a photonic NPU if you have one and the vendor toolkit installed.