Production-grade Kolmogorov-Arnold Networks โ TensorFlow + PyTorch + ONNX.
Project description
๐ kanx
Production-grade Kolmogorov-Arnold Networks
TensorFlow + PyTorch + ONNX โ one library, four surfaces.
๐ Why kanx?
pip install kanx265ร lower MSE than an MLP with 5ร fewer parameters
kanx is the first KAN library built for production systems โ not just research experiments.
โ One unified API โ Dual backend (TensorFlow + PyTorch) โ Real ONNX export (both backends) โ Built-in API serving โ Docker + Kubernetes ready
Train โ Export โ Serve โ Scale โ in one library.
โก 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โฆ]])
๐ง Same simplicity, real control
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])
๐ฅ PyTorch? Same API.
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]])
๐ฆ Installation
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)
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:
bash scripts/benchmark.sh
๐ง 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")
โ Dynamic batch โ Verified numerical consistency (1e-5) โ Works with ONNX Runtime / TensorRT / OpenVINO
๐ณ 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 |
๐ Research Paper
If you use kanx in academic work, please cite both the original paper and the library. Our work is formally documented and available as a preprint:
- ๐ Title: Bridging Theory and Practice with KANX
- ๐ DOI: https://doi.org/10.5281/zenodo.20430883
- ๐ Zenodo: https://zenodo.org/records/20430883
- ๐ Read Paper (preprint)
Citation
@article{mattral2026kanx,
title={Bridging Theory and Practice with KANX},
author={Myet, Min Htet},
year={2026},
doi={10.5281/zenodo.20430883},
publisher={Zenodo}
}
@article{liu2024kan,
title = {KAN: Kolmogorov-Arnold Networks},
author = {Liu, Ziming and Wang, Yixuan and Vaidya, Sachin and Ruehle,
Fabian and Halverson, James and Soljaฤiฤ, 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.5.tar.gz.
File metadata
- Download URL: kanx-0.1.5.tar.gz
- Upload date:
- Size: 41.7 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.13.9
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
4f0a3088bcce39fd3a014b5b59859f44dbaf23950f14866fb61e5d81701ea9bb
|
|
| MD5 |
f56baefc8c54b7a4947eec3749e08a88
|
|
| BLAKE2b-256 |
cca3746ee0227f0eb53ffefe501758e2df06d8ad283eaf49638d8cf86f26c478
|
File details
Details for the file kanx-0.1.5-py3-none-any.whl.
File metadata
- Download URL: kanx-0.1.5-py3-none-any.whl
- Upload date:
- Size: 34.0 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.13.9
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
4f72b0244cfdcc53b4aea0ead70e4bd9c330228ef4d6a5e0370523a982decafa
|
|
| MD5 |
5d7ef5710e9c4047335a8ef1d151d55d
|
|
| BLAKE2b-256 |
a6875b6c2a33675f860938209741c597a03fb746253497aea3e81edc168ee951
|