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

CI runs a CPU smoke gate that does not rewrite the published tables:

python benchmarks/compare_mlp.py --smoke
python benchmarks/real_world.py --smoke   # needs scikit-learn for California housing

python benchmarks/real_world.py --device cuda trains the PyTorch KAN and MLP in a child process, so TensorFlow and PyTorch do not initialize CUDA in one process. CPU runs and --smoke stay TensorFlow-only.

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

Drop edges that barely contribute. threshold is an absolute cutoff on the edge score. relative drops edges at or below that fraction of the strongest edge in the same layer. Pass samples to score each edge by its RMS on real inputs. The same method is on the PyTorch KAN and on MatrixKAN.

model.prune(relative=0.01)                  # keep edges above 1% of the peak
model.prune(threshold=1e-3, samples=X_train)

Plot what is left. Each edge is a learned 1-D function (SiLU residual plus B-spline). plot_edges needs matplotlib (pip install "kanx[viz]"). write_edge_viewer writes a single HTML file with a layer picker and a hover readout.

from kanx.viz import plot_edges, write_edge_viewer

plot_edges(model, path="edges.png")
write_edge_viewer(model, "edges.html")

Layer 0 of a KAN([2, 1]) fit to sin(π x₁) + 0.6 x₂². The first edge carries the sine. The second carries the quadratic.

The TensorFlow spline contraction can run in reduced precision, and the forward pass can be XLA-compiled. Weights and the returned tensor stay float32. The Cox-de Boor basis stays float32 as well: knot differences underflow in float16.

model = KAN([8, 16, 1], precision="bfloat16", jit_compile=True)

precision is float32 (default), float16, or bfloat16. jit_compile=True compiles on the first call. Batch size may change; the feature width may not. The same arguments exist on KANLinear. PyTorch is unchanged. Do not export a jit_compile=True layer to ONNX; trace the default float32 layer instead. Reduced precision is close to float32, not bit-exact. bfloat16 is the practical training dtype.

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

TensorBoard logging and symbolic edge fitting are included:

python -m kanx train --config configs/default.yaml --tensorboard
tensorboard --logdir logs/kanx

import torch
from kanx.torch import KAN, SymbolicFitter

model = KAN([1, 1])
x_1d = torch.linspace(-1, 1, 64)
name, r2 = SymbolicFitter(model).fit_edge(0, 0, 0, x_1d)

📦 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[viz]"         # +matplotlib for plot_edges
pip install "kanx[all]"         # everything (api + torch + onnx + hub + symbolic + viz + 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.
• kanx[viz] adds matplotlib so kanx.viz.plot_edges can draw the learned curves. The HTML viewer does not need it.

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

🏗️ Production Serving

Install kanx[api] and start the FastAPI app with uvicorn. Prometheus scrapes /metrics.

uvicorn api.app:app --port 8000

Kubernetes manifests (Deployment, Service, Ingress, HPA, PVC) live in k8s/. The same objects are parameterized in the Helm chart charts/kanx.

API Contract:

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

🌐 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
helm install kanx-api charts/kanx     # same objects, parameterized

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. Use the raw manifests or the chart, not both, when the release name is kanx-api. Install notes: charts/kanx/README.md.

🛠️ 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

• Unit, integration, live API, property-based, and performance tests in tests/
• 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: Python 3.10 / 3.11 / 3.12, lint, Docker smoke, MkDocs, and a CPU benchmark smoke gate
• CPU run of pytest tests/ --ignore=tests/test_e2e_live.py: 162 passed, 4 skipped (Helm render checks when helm is not installed, and MatrixKAN CUDA checks). Branch coverage of src/kanx on that run is 77% (layers.py 95%, model.py 91%). The older 94% / 113-test figures described a smaller suite. The distributed smoke in that count runs MirroredStrategy on two logical CPUs and torch.distributed with two gloo processes when no GPU is present.

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!