🚀 kanx

Production-grade Kolmogorov-Arnold Networks
TensorFlow + PyTorch + ONNX — one library, four surfaces.
> `kanx` is the only Kolmogorov-Arnold Network (KAN) library purpose-built for production deployment.

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. Prometheus metrics, TensorBoard logging, Hub and symbolic extras are now implemented.

⭐ Why kanx?

Most of the other KAN implementations are strictly academic. KANX bridges the gap between theory and reality by providing:

1. Multi-Backend Support: Native TensorFlow and PyTorch (MatrixKAN) implementations.
2. Real Deployment: True ONNX export, FastAPI (/api/predict), Docker, and Kubernetes configs out-of-the-box.
3. High-Throughput Options: Vectorized MatrixKAN to replace recursive B-splines with batched GEMMs for GPU acceleration.
4. Grid Calibration: Native adaptive and static grid fitting to prevent out-of-bounds input collapse.

📊 Benchmarks (reproducible, fair, multi-baseline)

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

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

⚠️ Grid calibration — two methods

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. Fix it one of two ways:

Static approach (pre-training):

from kanx import KAN, fit_grid_to_data
model = KAN([n_features, 64, 1])
fit_grid_to_data(model, X_train)              # one-time grid fit
model.fit(X_train, y_train, epochs=30)

Adaptive approach (during training — recommended):

model = KAN([n_features, 64, 1])
model.fit(X_train, y_train, epochs=15)
model.update_grid_from_samples(X_train)       # ← refine grid based on data
model.fit(X_train, y_train, epochs=15)        # continue training

kanx.check_input_range(model, X) will log a warning at inference if input exceeds the grid.

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

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

⚡ GPU-optimized MatrixKAN

For higher throughput on accelerators, use the vectorized MatrixKAN (replaces recursion with batched GEMM):

from kanx.torch import MatrixKAN

model = MatrixKAN([4, 32, 1])  # same interface as KAN
model.fit(X, y, epochs=30)      # ~1.5–2× faster on GPU vs standard KAN

📦 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[hub]"         # +HuggingFace Hub integration
pip install "kanx[symbolic]"    # +Symbolic regression hooks
pip install "kanx[all]"         # everything (api + torch + onnx + hub + symbolic + dev + docs)

Optional extras:

• kanx[api] adds FastAPI serving with /metrics Prometheus scraping.
• kanx[torch] adds the PyTorch backend, MatrixKAN, and symbolic helpers.
• kanx[hub] adds push_to_hub() / from_pretrained() for HuggingFace integration.
• kanx[symbolic] adds SymbolicFitter for post-hoc edge function extraction.

→ Open in Colab: Train a KAN in 2-to-5 minutes

🏗️ Production Serving

We include out-of-the-box serving. Simply install kanx[api] and run:

# Starts a FastAPI server with Prometheus scraping at /metrics
python -m kanx.serve 

API Contract:

• GET /api/health - Liveness & model load source
• GET /api/info - TF/Torch backend version and summary
• POST /api/predict - Batched inference

For enterprise scaling, see our /k8s directory for Helm charts and Kubernetes manifests.

🌐 REST API

docker run --rm -p 8000:8000 ghcr.io/mattral/kanx:latest
# or
uvicorn api.app:app --port 8000

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

⭐ Quality

• 95 tests across 8 files — unit, integration, E2E, property-based, performance regression
• 94% library coverage (99% layers, 100% model)
• Hypothesis property tests: partition of unity, shape invariants, gradient finiteness
• Numerical contracts: ONNX parity within 1e-5, save/load roundtrip identity
• Performance regression alarms: latency budgets on forward pass and predict
• CI matrix: Python 3.10 / 3.11 / 3.12 + lint + Docker smoke + MkDocs build

pytest tests/ -v --cov=src/kanx

📚 Documentation

📄 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:

• 📘 Title: KANX: A Production-Grade Open-Source Library for Kolmogorov-Arnold Networks
• 📍 DOI: https://doi.org/10.5281/zenodo.20615396
• 📂 Zenodo: https://zenodo.org/records/20430883
• 📄 Read Paper (preprint)
• 📄 Read Paper (ArXiv)

Citation

@article{mattral2026kanx,
  title={KANX: A Production-Grade Open-Source Library for Kolmogorov-Arnold Networks},
  author={Myet, Min Htet},
  year={2026},
  doi={10.5281/zenodo.20615396},
  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

• Liu et al., KAN: Kolmogorov-Arnold Networks — arXiv:2404.19756
• The Kolmogorov-Arnold representation theorem (Wikipedia)
• B-splines & de Boor algorithm — Carl de Boor (1972)

🤝 Contributing

PRs welcome! See CONTRIBUTING.md. Good places to start:

• 🔖 Good first issues
• 🗺️ roadmap.md — P0 / P1 / P2 backlog
• 💬 Discussions

📜 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!