KANX

🚀 kanx

Production-grade Kolmogorov-Arnold Networks
TensorFlow + PyTorch + ONNX — one library, four surfaces.

PyPI Downloads CI Python Docs Colab License DOI

KAN vs MLP benchmark

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])

🔥 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, 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")

✔ 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:

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

🤝 Contributing

PRs welcome! See CONTRIBUTING.md. Good places to start:

📜 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.

This site is open source. Improve this page.