Skip to content

SakunosukeCode/KANX-test

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

25 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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

About

Production-grade Kolmogorov-Arnold Networks TensorFlow + PyTorch + ONNX — one library, four surfaces. || 265× lower MSE than an MLP with 5× fewer parameters. One library. Two backends. Real ONNX export. Docker + Kubernetes ready.

Resources

License

Contributing

Security policy

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages

  • Python 93.7%
  • Jupyter Notebook 4.9%
  • Other 1.4%