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 DOI

KAN vs MLP benchmark


๐Ÿš€ Why kanx?

pip install kanx 265ร— lower MSE than an MLP with 5ร— fewer parameters

kanx is the first KAN library built for production systems โ€” not just research experiments.

โœ” One unified API โœ” Dual backend (TensorFlow + PyTorch) โœ” Real ONNX export (both backends) โœ” Built-in API serving โœ” Docker + Kubernetes ready

Train โ†’ Export โ†’ Serve โ†’ Scale โ€” in one library.


โšก 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โ€ฆ]])

๐Ÿง  Same simplicity, real control

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)

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:

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

โœ” 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.

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.5.tar.gz (41.7 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.5-py3-none-any.whl (34.0 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: kanx-0.1.5.tar.gz
  • Upload date:
  • Size: 41.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.9

File hashes

Hashes for kanx-0.1.5.tar.gz
Algorithm Hash digest
SHA256 4f0a3088bcce39fd3a014b5b59859f44dbaf23950f14866fb61e5d81701ea9bb
MD5 f56baefc8c54b7a4947eec3749e08a88
BLAKE2b-256 cca3746ee0227f0eb53ffefe501758e2df06d8ad283eaf49638d8cf86f26c478

See more details on using hashes here.

File details

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

File metadata

  • Download URL: kanx-0.1.5-py3-none-any.whl
  • Upload date:
  • Size: 34.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.9

File hashes

Hashes for kanx-0.1.5-py3-none-any.whl
Algorithm Hash digest
SHA256 4f72b0244cfdcc53b4aea0ead70e4bd9c330228ef4d6a5e0370523a982decafa
MD5 5d7ef5710e9c4047335a8ef1d151d55d
BLAKE2b-256 a6875b6c2a33675f860938209741c597a03fb746253497aea3e81edc168ee951

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