Production-grade Kolmogorov-Arnold Networks — TensorFlow + PyTorch + ONNX.
Project description
🚀 kanx
Production-grade Kolmogorov-Arnold Networks
TensorFlow + PyTorch + ONNX — one library, four surfaces.
pip install kanx· A small KAN beats a 10× larger MLP on smooth, separable targets — honest, param-matched benchmark below. One library. Two backends. Real ONNX export. Docker + Kubernetes ready.
⚡ The 30-second magic moment
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…]])
⚠️ The #1 KAN production gotcha — calibrate your grid to your data
KANs use B-splines on a fixed input range (default
[-1, 1]). If your inputs fall outside that range, the spline path silently returns zero and you only get the SiLU residual — degraded accuracy with no error message. Fix it in one line:from kanx import KAN, fit_grid_to_data model = KAN([n_features, 64, 1]) fit_grid_to_data(model, X_train) # ← critical for real data model.fit(X_train, y_train, epochs=30)
kanx.check_input_range(model, X)will log a warning at inference if input exceeds the grid. The Roadmap tracks adaptive in-training grid updates as the v0.2 P0 item.
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]])
📦 Install
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
📊 Benchmarks (reproducible, fair, multi-baseline)
Synthetic 2-D regression target y = sin(π·x₁) + cos(2π·x₂),
100 epochs, Adam(lr=1e-2), batch=128, CPU.
| Model | Params | Train (s) | Infer 4k (ms) | Test MSE |
|---|---|---|---|---|
| KAN[2,16,1] | 432 | 12.50 | 68.64 | 2.14 × 10⁻⁵ |
| KAN[2,32,1] | 864 | 16.62 | 25.52 | 4.44 × 10⁻⁴ |
| MLP[2,32,1] | 129 | 5.07 | 6.17 | 4.61 × 10⁻¹ (undersized) |
| MLP[2,16,16,1] | 337 | 5.46 | 4.08 | 1.60 × 10⁻³ |
| MLP[2,64,64,1] | 4 417 | 6.00 | 5.74 | 5.51 × 10⁻⁴ |
Honest read. The smallest KAN (432 params) wins on this smooth separable target. The same KAN is ~10–15× slower at inference than a same-MSE MLP because each edge does a B-spline evaluation. On non-smooth or high-dimensional targets, this picture often reverses. We do not claim KANs are universally better than MLPs.
Reproduce with python benchmarks/compare_mlp.py (quick, 100 epochs) or
python benchmarks/compare_mlp.py --long (1000 epochs + early-stopping).
🧠 How kanx compares to other KAN libraries
| 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.
🌐 REST API
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.
🔄 ONNX export
# 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 / ☸️ Kubernetes
docker run --rm -p 8000:8000 ghcr.io/mattral/kanx:latest
kubectl apply -f k8s/ # Deployment + Service + Ingress + HPA + PVC
K8s 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.
🛠️ CLI
python -m kanx info # versions
python -m kanx train --config configs/default.yaml # train
python -m kanx predict --checkpoint model.keras --input X.json
📚 Documentation
→ 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 |
🧬 Citing kanx & the KAN paper
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}
}
References
- 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)
🤝 Contributing
PRs welcome! See CONTRIBUTING.md. Good places to start:
- 🔖 Good first issues
- 🗺️
roadmap.md— P0 / P1 / P2 backlog - 💬 Discussions
📜 License
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.
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file kanx-0.1.6.tar.gz.
File metadata
- Download URL: kanx-0.1.6.tar.gz
- Upload date:
- Size: 45.6 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.15
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
b639930d5c14ce45ee5aa89fb7c9570561b863fc52a0e69b28336a52c157f55a
|
|
| MD5 |
cb18bfa93f5a4d3fbb344629c4a7df86
|
|
| BLAKE2b-256 |
aa3cbae7c5c7ecee1bfca18f0ae1adc29f742b5a47aac7b0d7bb6f7ecb1677ba
|
File details
Details for the file kanx-0.1.6-py3-none-any.whl.
File metadata
- Download URL: kanx-0.1.6-py3-none-any.whl
- Upload date:
- Size: 35.9 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.15
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
0187020b1ab5b4e4434ff93d508e35034237b79e104dde5de2f67158e26026ce
|
|
| MD5 |
6233a81ccbbcc78c96d012fc99822f4b
|
|
| BLAKE2b-256 |
6b6ae39ca1fffd8f4b7c3158ff49768ae9bcec699a1b3d2727ff646bc329d9ba
|