Skip to main content

Production-grade Kolmogorov-Arnold Networks — TensorFlow + PyTorch + ONNX.

Project description

🚀 kanx

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

PyPI Downloads CI Python Docs Colab License

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

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


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

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

kanx-0.1.6.tar.gz (45.6 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

kanx-0.1.6-py3-none-any.whl (35.9 kB view details)

Uploaded Python 3

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

Hashes for kanx-0.1.6.tar.gz
Algorithm Hash digest
SHA256 b639930d5c14ce45ee5aa89fb7c9570561b863fc52a0e69b28336a52c157f55a
MD5 cb18bfa93f5a4d3fbb344629c4a7df86
BLAKE2b-256 aa3cbae7c5c7ecee1bfca18f0ae1adc29f742b5a47aac7b0d7bb6f7ecb1677ba

See more details on using hashes here.

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

Hashes for kanx-0.1.6-py3-none-any.whl
Algorithm Hash digest
SHA256 0187020b1ab5b4e4434ff93d508e35034237b79e104dde5de2f67158e26026ce
MD5 6233a81ccbbcc78c96d012fc99822f4b
BLAKE2b-256 6b6ae39ca1fffd8f4b7c3158ff49768ae9bcec699a1b3d2727ff646bc329d9ba

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page