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/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.4.tar.gz (41.0 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.4-py3-none-any.whl (33.6 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: kanx-0.1.4.tar.gz
  • Upload date:
  • Size: 41.0 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.4.tar.gz
Algorithm Hash digest
SHA256 d9917d05ae0c2dac273e61c6b7591663152bc9ed9acc5ac94cbfbe8719aab125
MD5 f0f2d095cabd68f1a02c9fca7077eac4
BLAKE2b-256 ddbe1d93c73b1f3a24c8840909d59bf40ecda7c5d6bdc1341061d5bbb7427dc9

See more details on using hashes here.

File details

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

File metadata

  • Download URL: kanx-0.1.4-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.4-py3-none-any.whl
Algorithm Hash digest
SHA256 a07409ab1e01e8c9b74528a80f1b6cca94c972523f1953c71d6a1c8535f6824c
MD5 1a69ad1af2ecdb8e73147e71ef688126
BLAKE2b-256 d4e7bfc731f8dc9f8ad1d0587dce921d3fc33af12eae3346267de373f2e31818

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