Documentation · Quickstart · Tutorials · Kernel benchmarks · Roadmap · Changelog

A Kolmogorov–Arnold Network replaces every weight of a neural network with a small learnable function. kanforge gives you those layers with 26 families of edge functions, on PyTorch, Apple MLX and JAX — one parameter layout everywhere, fused GPU kernels, tools that turn trained curves back into formulas, and exporters that take a model all the way down to an FPGA.

pip install "kanforge[torch] @ git+https://github.com/fx-xf/kanforge"

Combine extras as you need them: torch, mlx (Apple silicon), jax (Flax NNX, Python ≥ 3.11), symbolic, plot, hf, export (ONNX), coreml. The core depends on NumPy only.

Backend	Import	Devices
PyTorch	kanforge.torch	CUDA · ROCm · Apple MPS · CPU
MLX	kanforge.mlx	Apple silicon GPU / CPU
JAX · Flax NNX	kanforge.jax	GPU · TPU · CPU
NumPy, float64	kanforge.reference	CPU — the reference every test checks against

A KAN in three lines. Any basis, any backend, the same parameters.

import torch
from kanforge.torch import KAN

model = KAN([2, 5, 1], basis="bspline")          # or "chebyshev", "fourier", "morlet", …
y = model(torch.rand(64, 2))

Layers with character. Exact derivatives of any order, the group-rational layer of KAT, multiplication nodes of KAN 2.0.

import kanforge as kf
from kanforge.torch import GRKAN, KANLinear, MultKAN

layer = KANLinear(8, 4, kf.Jacobi(6, alpha=0.5, beta=0.5), edge_scale=True)
d2 = layer.edge_derivative(torch.rand(32, 8), order=2)   # exact d²y/dx², (32, 4, 8)
ffn = GRKAN(768, 3072, groups=8)                         # KAT's group-rational FFN
model = MultKAN([2, (2, 1), 1])                          # two sum nodes, one product node

Grids that follow your data. Adaptive per-input grids, refinement that keeps the function, conversion between bases, and MLPs turned into KANs without changing a single output.

from kanforge.torch import KAN, from_mlp, regularization

model = KAN([2, 5, 1], kf.BSpline(5, 3), coef_init="lecun")
x = torch.rand(1024, 2) * 4 - 2
model.update_grid(x)                          # per-input adaptive grids
model.refine_grid(10)                         # finer grid, same function
model[0].change_basis(kf.Chebyshev(12))       # convert a layer to another basis
loss = model(x).pow(2).mean() + 1e-3 * regularization(model, x)

mlp = torch.nn.Sequential(torch.nn.Linear(4, 16), torch.nn.GELU(), torch.nn.Linear(16, 1))
kan = from_mlp(mlp)                           # equals the MLP exactly — now fine-tune the splines

From curves to formulas — and back.

import sympy
from kanforge import interpret as ki

pruned = ki.prune(model, x)                   # drop unimportant nodes and edges
ki.auto_symbolic(pruned, x)                   # fit c·f(a·x + b) + d to every edge
print(ki.symbolic_formula(pruned, ["x", "y"]))

X, Y = sympy.symbols("x y")
kan, _ = ki.kanpile(sympy.exp(sympy.sin(sympy.pi * X) + Y**2), [X, Y])   # formula → MultKAN

Bigger models. KAN convolutions, transformer surgery, derivatives for physics-informed training.

import transformers
from kanforge.torch import KANConv2d, kanify, derivative

gpt2 = transformers.AutoModelForCausalLM.from_pretrained("gpt2")
kanify(gpt2)                                  # FFN projections → KANLinear, identical outputs
conv = KANConv2d(3, 16, 3, padding=1, basis="chebyshev")
u_xx = derivative(net, x, order=2)            # PDE residuals for PINNs

Anywhere it needs to run.

import kanforge
from kanforge.export import fpga
from kanforge.export.onnx_export import export_onnx

kanforge.save(model, "model.npz")                         # from any backend …
jax_model = kanforge.load("model.npz", backend="jax")     # … into any other
mlx_model = kanforge.convert.convert(model, "mlx")
export_onnx(model, "model.onnx", x, size=256, bits=8)     # int8 lookup tables
fpga.export_fpga(model, "rtl/", x, address_bits=8, value_bits=8)   # Verilog + bit-exact model

Every curve above is computed by kanforge.reference. Every basis provides derivatives of any order in closed form, shares its trainable parameters per input, per group or per layer, and runs through the same fused kernels.

KANs are not a free win. At equal FLOPs they are 5–50× more accurate on smooth, compositional problems — physical laws, symbolic formulas, scientific tabular data — and slightly behind MLPs on raw pixels. The full tables are in benchmarks/fair_results.md; kernel timings on Apple M5 and RTX 3060 are in BENCHMARKS.md.

Guide	bases · layers and backends · training · fused kernels · interpretation · architecture blocks · JAX · deployment
Tutorials	regression · formula discovery · PINNs · operator learning · images · transformers · graphs · time series · neural ODEs · reinforcement learning · deployment · backends
API reference	generated from the docstrings
Examples	the twelve tutorials as runnable scripts

Issues and pull requests are welcome — read CONTRIBUTING and the roadmap first. kanforge is released under the Apache-2.0 licence; ideas and code from other KAN projects are credited in NOTICE. If kanforge helps your research, please cite it (CITATION.cff) together with the original KAN paper (arXiv:2404.19756).

_{The figures in this README are drawn by assets/src/build.py — every curve in them is computed by kanforge.}