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  ·  265× lower MSE than an MLP with 5× fewer parameters. 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…]])

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)

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.


🧠 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/Kolmogorov-Arnold-Networks/ (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/Kolmogorov-Arnold-Networks}}
}

@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.3.tar.gz (41.1 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.3-py3-none-any.whl (33.6 kB view details)

Uploaded Python 3

File details

Details for the file kanx-0.1.3.tar.gz.

File metadata

  • Download URL: kanx-0.1.3.tar.gz
  • Upload date:
  • Size: 41.1 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.3.tar.gz
Algorithm Hash digest
SHA256 8adc05086f7879aa4ba9f18906243b9d6681ccd77b6c50b9bbb2d3ce4c96526e
MD5 5630b320759b90db6b66dfc2f440147a
BLAKE2b-256 b028b2210bfe70a4bf74ac1b25f47a003b70b676a0dff2e9b1ae6339aba034fe

See more details on using hashes here.

File details

Details for the file kanx-0.1.3-py3-none-any.whl.

File metadata

  • Download URL: kanx-0.1.3-py3-none-any.whl
  • Upload date:
  • Size: 33.6 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.3-py3-none-any.whl
Algorithm Hash digest
SHA256 9743a7d791874c5b56fe969bacb3f624905576a82200632a90e485d7f9140fd2
MD5 63000fa5c91d7d8c77f03b76f924aacd
BLAKE2b-256 0e0a247eb25c85b3744a118d22562ab2730ea594351960f9ecd9cab5376e6c7e

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