kanx 0.1.0

Production-grade Kolmogorov-Arnold Networks — TensorFlow + PyTorch + ONNX.

🚀 kanx

Production-grade Kolmogorov-Arnold Networks — TensorFlow + PyTorch + ONNX

265× lower test MSE than an MLP with 5× fewer parameters. One library. Two backends. Real ONNX export. Docker + Kubernetes ready.

---

⚡ The 10-line magic moment

=== "🔷 TensorFlow (default)"

```python
from kanx import KAN, train, load_config
import numpy as np

X = np.random.uniform(-1, 1, (512, 2)).astype("float32")
y = np.sin(np.pi * X[:, :1])

model, _ = train(load_config("configs/default.yaml"), X, y)
model.predict(X[:3])
```

=== "🔶 PyTorch"

```python
from kanx.torch import KAN, Trainer
import torch

model = KAN([2, 64, 1])
X = torch.randn(512, 2); y = torch.sin(torch.pi * X[:, :1])

Trainer(model).fit(X, y, epochs=30, lr=1e-2)
model.predict([[0.5, 0.2]])
```

=== "🌐 REST API"

```bash
docker run --rm -p 8000:8000 ghcr.io/mattral/kanx:latest
curl -X POST http://localhost:8000/api/predict \
     -H 'content-type: application/json' \
     -d '{"x": [[0.1, -0.2], [0.5, 0.7]]}'
```

=== "🔄 ONNX export"

```python
from kanx.torch import KAN, export_onnx
import torch

model = KAN([2, 64, 1])
export_onnx(model, "kan.onnx", sample_input=torch.zeros(1, 2))
# → run anywhere with ONNX Runtime, TensorRT, OpenVINO, CoreML, …
```

---

🧠 Why kanx?

	Research-y KAN repos	kanx
Two backends (TF + Torch)	❌	✅
Vectorized B-spline	⚠️ Python loops	✅ pure einsum
ONNX export	❌	✅ TF & Torch
FastAPI service	❌	✅ hot-swap, 5 endpoints
Docker + K8s + HPA	❌	✅
CI/CD + release pipeline	❌	✅ PyPI + GHCR + Pages
Property-based tests	❌	✅ Hypothesis
Numerical-contract tests	❌	✅ partition of unity
Test coverage	~30%	94%
Performance budgets	❌	✅ regression alarms

---

📊 Benchmarks

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.

---

📦 Install

pip install kanx              # core (TensorFlow)
pip install "kanx[torch]"     # +PyTorch backend
pip install "kanx[api]"       # +FastAPI service
pip install "kanx[all]"       # everything

Or development install:

git clone https://github.com/Mattral/Kolmogorov-Arnold-Networks
cd Kolmogorov-Arnold-Networks
pip install -e ".[dev,api,torch]"
bash scripts/test.sh          # 95 tests, 94% coverage

---

🧪 What's in the box

kanx/
├── src/kanx/             # library
│   ├── layers.py         # TF KANLinear + B-spline primitives
│   ├── model.py          # TF KAN sequential
│   ├── train.py          # train() + optim/loss factories
│   ├── inference.py      # predict + save/load
│   ├── onnx_export.py    # TF → ONNX (via tf2onnx)
│   ├── config.py         # YAML config + strict validator
│   ├── utils.py          # logging, seeding
│   ├── __main__.py       # CLI: `python -m kanx`
│   └── torch/            # ⚡ PyTorch backend (parallel surface)
│       ├── layers.py     # Torch KANLinear + B-spline
│       ├── model.py      # Torch KAN + save/load
│       ├── trainer.py    # Trainer.fit(...)
│       └── onnx_export.py # Torch → ONNX
├── api/                  # FastAPI service
├── tests/                # 95 tests — unit + property + perf + E2E
├── benchmarks/           # KAN vs MLP harness
├── examples/             # quickstart, function_fit, MNIST
├── configs/              # YAML training configs
├── notebooks/            # Colab-ready notebooks
├── k8s/                  # Deployment + Service + Ingress + HPA + PVC
├── docs/                 # MkDocs Material site
└── .github/workflows/    # CI + release (PyPI + GHCR + Pages)

---

🌐 REST API

uvicorn api.app:app --port 8000           # or docker compose up

Method	Path	Purpose
GET	/api/health	Liveness + model load source
GET	/api/info	Version + TF version + model summary
POST	/api/predict	Inference (single or batch)
POST	/api/load	Hot-swap checkpoint
POST	/api/reset	Re-init from KANX_CONFIG

All endpoints validate at the boundary: wrong feature count → 400, oversized batch → 413, missing checkpoint → 404. The startup contract loads KANX_CHECKPOINT if it exists, otherwise falls back to a fresh model built from KANX_CONFIG.

---

🐳 Docker / ☸️ Kubernetes

docker run --rm -p 8000:8000 ghcr.io/mattral/kanx:latest
kubectl apply -f k8s/        # Deployment + Service + Ingress + HPA + PVC

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

Full docs are published as a MkDocs Material site:

→ https://mattral.github.io/Kolmogorov-Arnold-Networks/

Or read the source files under documentations/.

---

🧬 Research positioning

KANs were introduced by Liu et al. (arXiv:2404.19756). This repo is not a reproduction of the paper — it's an opinionated engineering layer that takes the core idea (learnable spline activations on edges, SiLU residual) and packages it for real-world MLOps:

• Two backends so you can drop kanx into either TF or Torch stacks.
• ONNX so you can deploy anywhere.
• A REST API + Docker so you can serve.
• A test suite that catches bugs before users do.

---

🤝 Contributing

PRs welcome. Please:

1. Open an issue first for any non-trivial change.
2. Run bash scripts/test.sh — it must stay green.
3. Add a test for any new behaviour (we're at 94% coverage; let's keep climbing).
4. Update roadmap.md when shipping a feature.

---

📜 License

Apache 2.0. Use it. Ship it. Tell us when you do.

Built with care so KANs are actually usable, not just academically interesting.