Production-grade Kolmogorov-Arnold Networks
TensorFlow + PyTorch + ONNX — one library, four surfaces.
pip install kanx· 265× lower MSE than an MLP with 5× fewer parameters. One library. Two backends. Real ONNX export. Docker + Kubernetes ready.
import kanx
# Build, train, predict — in one call. No config files. No compile dance.
model = kanx.quickstart() # trains on synthetic 2-D data
model.predict([[0.5, 0.2]]) # → array([[1.04…]])Want more control? Same simplicity, your data:
from kanx import KAN
import numpy as np
X = np.random.uniform(-1, 1, (1024, 2)).astype("float32")
y = np.sin(np.pi * X[:, :1]) + X[:, 1:2] ** 2
model = KAN([2, 64, 1])
model.fit(X, y, epochs=30, verbose=0) # auto-compiles with Adam+MSE
model.predict(X[:3])Prefer PyTorch?
from kanx.torch import KAN
import torch
model = KAN([2, 64, 1])
X = torch.randn(1024, 2); y = torch.sin(torch.pi * X[:, :1])
model.fit(X, y, epochs=30, lr=1e-2) # one-liner, same semantics
model.predict([[0.5, 0.2]])pip install kanx # core (TensorFlow)
pip install "kanx[torch]" # +PyTorch backend
pip install "kanx[onnx]" # +tf2onnx + onnxruntime
pip install "kanx[api]" # +FastAPI service
pip install "kanx[all]" # everything (api + torch + onnx + dev + docs)→ Open in Colab: Train a KAN in 2 minutes
Synthetic 2-D regression target y = sin(π·x₁) + cos(2π·x₂),
30 epochs, Adam(lr=1e-2), batch=128, CPU.
| Model | Params | Train (s) | Infer 4k (ms) | Test MSE |
|---|---|---|---|---|
| KAN[2,16,1] | 432 | 4.18 | 35.71 | 6.4 × 10⁻⁵ |
| KAN[2,32,1] | 864 | 5.31 | 28.50 | 1.7 × 10⁻⁵ |
| MLP[2,64,64,1] | 4 417 | 2.04 | 6.88 | 4.5 × 10⁻³ |
Reproduce with bash scripts/benchmark.sh.
| pykan | efficient-kan | mlx-kan | kanx | |
|---|---|---|---|---|
| Framework | PyTorch | PyTorch | MLX (Apple Silicon) | TF + PyTorch |
| Vectorized B-spline | partial | ✅ | ✅ | ✅ |
| ONNX export | ❌ | ❌ | ❌ | ✅ both backends |
| REST API service | ❌ | ❌ | ❌ | ✅ FastAPI |
| Docker + K8s | ❌ | ❌ | ❌ | ✅ |
| Property-based tests | ❌ | ❌ | ❌ | ✅ Hypothesis |
| Test coverage | research | research | research | 92% |
| PyPI | ✅ | ✅ | ✅ | ✅ |
| CI/CD release pipeline | ❌ | ❌ | ❌ | ✅ PyPI + GHCR + Pages |
kanx is the only KAN library purpose-built for production deployment.
Research-y libs are great for novel experiments; kanx is what you ship.
docker run --rm -p 8000:8000 ghcr.io/mattral/kanx:latest
# or
uvicorn api.app:app --port 8000| Method | Path | Purpose |
|---|---|---|
GET |
/api/health |
Liveness + model load source |
GET |
/api/info |
Version + TF/Torch + model summary |
POST |
/api/predict |
Inference (single or batch) |
POST |
/api/load |
Hot-swap checkpoint |
POST |
/api/reset |
Re-init from KANX_CONFIG |
curl -X POST http://localhost:8000/api/predict \
-H 'content-type: application/json' \
-d '{"x": [[0.1, -0.2], [0.5, 0.7]]}'The startup contract loads KANX_CHECKPOINT if it exists, otherwise falls
back to a fresh model built from KANX_CONFIG. Boundaries are validated:
wrong feature count → 400, oversized batch → 413, missing checkpoint → 404.
# From PyTorch
from kanx.torch import KAN, export_onnx
model = KAN([2, 64, 1])
export_onnx(model, "kan.onnx")
# From TensorFlow
from kanx import KAN, export_onnx_tf
import tensorflow as tf
model = KAN([2, 64, 1]); model(tf.zeros((1, 2)))
export_onnx_tf(model, "kan.onnx")Both exports include a dynamic batch axis and are tested to be numerically identical to the eager model within 1e-5. Run anywhere with ONNX Runtime, TensorRT, OpenVINO, CoreML, …
docker run --rm -p 8000:8000 ghcr.io/mattral/kanx:latest
kubectl apply -f k8s/ # Deployment + Service + Ingress + HPA + PVCK8s manifests ship with rolling updates, readiness/liveness probes on
/api/health, an HPA (2 ↔ 10 replicas, CPU-target 70%) and a PVC for the
model registry.
python -m kanx info # versions
python -m kanx train --config configs/default.yaml # train
python -m kanx predict --checkpoint model.keras --input X.json→ https://mattral.github.io/KANX/ (MkDocs Material)
| Page | What's inside |
|---|---|
| Quickstart | Train your first KAN in 60 seconds |
| Architecture | Package layout, module contracts |
| System Design | Serving topology, scaling, failure modes |
| REST API | Endpoint reference + curl examples |
| Testing | Test pyramid, numerical invariants |
| Deployment | CI/CD, rollout, observability |
| Benchmarks | KAN vs MLP — methodology + numbers |
If you use kanx in academic work, please cite both the original paper and the library:
@misc{kanx,
author = {kanx contributors},
title = {kanx: Production-grade Kolmogorov-Arnold Networks},
year = {2026},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/Mattral/KANX}}
}
@article{liu2024kan,
title = {KAN: Kolmogorov-Arnold Networks},
author = {Liu, Ziming and Wang, Yixuan and Vaidya, Sachin and Ruehle,
Fabian and Halverson, James and Solja{\v c}i{\'c}, Marin and
Hou, Thomas Y. and Tegmark, Max},
journal = {arXiv preprint arXiv:2404.19756},
year = {2024}
}- Liu et al., KAN: Kolmogorov-Arnold Networks — arXiv:2404.19756
- The Kolmogorov-Arnold representation theorem (Wikipedia)
- B-splines & de Boor algorithm — Carl de Boor (1972)
PRs welcome! See CONTRIBUTING.md. Good places to start:
- 🔖 Good first issues
- 🗺️
roadmap.md— P0 / P1 / P2 backlog - 💬 Discussions
Apache 2.0. Use it. Ship it. Tell us when you do — we'd love to hear how kanx is being used in the wild.
⭐ Star the repo if kanx saved you time.
