inkan 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- inkan-0.1.0/LICENSE +21 -0
- inkan-0.1.0/PKG-INFO +220 -0
- inkan-0.1.0/README.md +192 -0
- inkan-0.1.0/pyproject.toml +43 -0
- inkan-0.1.0/setup.cfg +4 -0
- inkan-0.1.0/src/inkan/__init__.py +24 -0
- inkan-0.1.0/src/inkan/basis.py +53 -0
- inkan-0.1.0/src/inkan/layer.py +112 -0
- inkan-0.1.0/src/inkan/network.py +50 -0
- inkan-0.1.0/src/inkan/visualize.py +309 -0
- inkan-0.1.0/src/inkan.egg-info/PKG-INFO +220 -0
- inkan-0.1.0/src/inkan.egg-info/SOURCES.txt +15 -0
- inkan-0.1.0/src/inkan.egg-info/dependency_links.txt +1 -0
- inkan-0.1.0/src/inkan.egg-info/requires.txt +6 -0
- inkan-0.1.0/src/inkan.egg-info/top_level.txt +1 -0
- inkan-0.1.0/tests/test_basis.py +151 -0
- inkan-0.1.0/tests/test_layer.py +96 -0
inkan-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Naveen Mysore
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
inkan-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,220 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: inkan
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Fast B-spline KAN layers — 6-15x faster than Cox-de Boor, exact B-spline values
|
|
5
|
+
Author: Naveen Mysore
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/naveenmysore/flashkan
|
|
8
|
+
Project-URL: Repository, https://github.com/naveenmysore/flashkan
|
|
9
|
+
Keywords: kan,kolmogorov-arnold,b-spline,neural-network,deep-learning
|
|
10
|
+
Classifier: Development Status :: 3 - Alpha
|
|
11
|
+
Classifier: Intended Audience :: Science/Research
|
|
12
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
13
|
+
Classifier: Programming Language :: Python :: 3
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
19
|
+
Requires-Python: >=3.9
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
License-File: LICENSE
|
|
22
|
+
Requires-Dist: torch>=2.0
|
|
23
|
+
Provides-Extra: dev
|
|
24
|
+
Requires-Dist: pytest>=7.0; extra == "dev"
|
|
25
|
+
Requires-Dist: torchvision>=0.15; extra == "dev"
|
|
26
|
+
Requires-Dist: matplotlib>=3.5; extra == "dev"
|
|
27
|
+
Dynamic: license-file
|
|
28
|
+
|
|
29
|
+
# FlashKAN
|
|
30
|
+
|
|
31
|
+
Fast B-spline [Kolmogorov-Arnold Network](https://arxiv.org/abs/2404.19756) layers for PyTorch.
|
|
32
|
+
|
|
33
|
+
**6-15x faster** than standard Cox-de Boor implementations (PyKAN, efficient-kan), **faster than Gaussian RBF** alternatives (FastKAN) — while producing **exact B-spline basis values** with compact support, C2 continuity, and partition of unity.
|
|
34
|
+
|
|
35
|
+
## How it works
|
|
36
|
+
|
|
37
|
+
Standard KAN implementations compute B-spline basis functions using the Cox-de Boor recursion — 3 sequential passes for cubic splines, each creating intermediate tensors. FlashKAN replaces this with the truncated power closed form:
|
|
38
|
+
|
|
39
|
+
```
|
|
40
|
+
N(u) = (1/6) [relu(u)³ - 4·relu(u-1)³ + 6·relu(u-2)³ - 4·relu(u-3)³ + relu(u-4)³]
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
This single expression computes exact B-spline values with no recursion, no span lookups, and no gather operations. `torch.compile` fuses all elementwise ops into one GPU kernel.
|
|
44
|
+
|
|
45
|
+
## Installation
|
|
46
|
+
|
|
47
|
+
```bash
|
|
48
|
+
pip install flashkan
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
**Requirements:** Python >= 3.9, PyTorch >= 2.0
|
|
52
|
+
|
|
53
|
+
**Supported devices:** CPU, CUDA (NVIDIA), MPS (Apple Silicon)
|
|
54
|
+
|
|
55
|
+
### From source
|
|
56
|
+
|
|
57
|
+
```bash
|
|
58
|
+
git clone https://github.com/NAVEENMN/flashkan.git
|
|
59
|
+
cd flashkan
|
|
60
|
+
pip install -e .
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
## Quick start
|
|
64
|
+
|
|
65
|
+
```python
|
|
66
|
+
import torch
|
|
67
|
+
from flashkan import KANLayer, KANNetwork
|
|
68
|
+
|
|
69
|
+
# Drop-in replacement for nn.Linear
|
|
70
|
+
layer = KANLayer(784, 64)
|
|
71
|
+
x = torch.randn(32, 784)
|
|
72
|
+
y = layer(x) # [32, 64]
|
|
73
|
+
|
|
74
|
+
# Multi-layer network
|
|
75
|
+
net = KANNetwork([784, 64, 10])
|
|
76
|
+
y = net(torch.randn(32, 784)) # [32, 10]
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
### MNIST example
|
|
80
|
+
|
|
81
|
+
```python
|
|
82
|
+
import torch
|
|
83
|
+
import torch.nn as nn
|
|
84
|
+
from flashkan import KANNetwork
|
|
85
|
+
|
|
86
|
+
model = KANNetwork([784, 64, 10])
|
|
87
|
+
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
|
|
88
|
+
criterion = nn.CrossEntropyLoss()
|
|
89
|
+
|
|
90
|
+
# Standard PyTorch training loop
|
|
91
|
+
for images, labels in train_loader:
|
|
92
|
+
output = model(images.view(-1, 784))
|
|
93
|
+
loss = criterion(output, labels)
|
|
94
|
+
optimizer.zero_grad()
|
|
95
|
+
loss.backward()
|
|
96
|
+
optimizer.step()
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
See [`examples/`](examples/) for complete runnable scripts.
|
|
100
|
+
|
|
101
|
+
## Visualization
|
|
102
|
+
|
|
103
|
+
FlashKAN includes built-in visualization for learned activation functions — similar to PyKAN's `model.plot()`.
|
|
104
|
+
|
|
105
|
+
```python
|
|
106
|
+
from flashkan import KANNetwork, plot_basis, plot_activations, plot_network
|
|
107
|
+
|
|
108
|
+
model = KANNetwork([784, 32, 10], grid_size=8)
|
|
109
|
+
# ... train on MNIST ...
|
|
110
|
+
|
|
111
|
+
plot_basis(model.layers[0]) # B-spline basis bumps
|
|
112
|
+
plot_activations(model.layers[1]) # learned curves per edge
|
|
113
|
+
plot_network(model) # full network diagram
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
### B-spline basis functions
|
|
117
|
+
|
|
118
|
+
The 8 basis bumps (grid_size=5, degree=3) — compact support, smooth overlap:
|
|
119
|
+
|
|
120
|
+

|
|
121
|
+
|
|
122
|
+
### Learned activation functions
|
|
123
|
+
|
|
124
|
+
After training on MNIST, each edge learns a unique activation curve.
|
|
125
|
+
Cyan = total, red dashed = spline component, green dotted = SiLU base:
|
|
126
|
+
|
|
127
|
+

|
|
128
|
+
|
|
129
|
+
### Network diagram
|
|
130
|
+
|
|
131
|
+
Full [784 → 32 → 10] network with learned curves on edges:
|
|
132
|
+
|
|
133
|
+

|
|
134
|
+
|
|
135
|
+
## API
|
|
136
|
+
|
|
137
|
+
### `KANLayer(in_features, out_features, grid_size=5, spline_order=3)`
|
|
138
|
+
|
|
139
|
+
A single KAN layer. Drop-in replacement for `nn.Linear`.
|
|
140
|
+
|
|
141
|
+
| Parameter | Default | Description |
|
|
142
|
+
|---|---|---|
|
|
143
|
+
| `in_features` | — | Input dimension |
|
|
144
|
+
| `out_features` | — | Output dimension |
|
|
145
|
+
| `grid_size` | 5 | Number of knot intervals (more = finer approximation) |
|
|
146
|
+
| `spline_order` | 3 | B-spline degree (3 = cubic, recommended) |
|
|
147
|
+
| `grid_range` | (-1, 1) | Input range for the spline grid |
|
|
148
|
+
|
|
149
|
+
### `KANNetwork(layer_dims, grid_size=5, spline_order=3)`
|
|
150
|
+
|
|
151
|
+
Stack of KAN layers.
|
|
152
|
+
|
|
153
|
+
```python
|
|
154
|
+
# 3-layer KAN: 784 -> 128 -> 64 -> 10
|
|
155
|
+
net = KANNetwork([784, 128, 64, 10])
|
|
156
|
+
```
|
|
157
|
+
|
|
158
|
+
## Benchmarks
|
|
159
|
+
|
|
160
|
+
Forward pass time (ms) on Apple M-series GPU (MPS), batch=256:
|
|
161
|
+
|
|
162
|
+
| Layer | MNIST (784→64) | FashionMNIST (784→64) | CIFAR-10 (3072→64) |
|
|
163
|
+
|---|---|---|---|
|
|
164
|
+
| **FlashKAN (compiled)** | **0.27** | **0.20** | **0.39** |
|
|
165
|
+
| FastKAN (Gaussian RBF) | 0.38 | 0.31 | 0.99 |
|
|
166
|
+
| NoGather (unrolled) | 0.99 | 1.02 | 4.85 |
|
|
167
|
+
| Vanilla (Cox-de Boor) | 1.69 | 1.67 | 5.98 |
|
|
168
|
+
|
|
169
|
+
FlashKAN is **6.3x faster** than vanilla Cox-de Boor on MNIST and **15.3x faster** on CIFAR-10.
|
|
170
|
+
|
|
171
|
+
### Why it's fast
|
|
172
|
+
|
|
173
|
+
91% of a standard KAN forward pass is spent computing B-spline basis functions. FlashKAN eliminates this bottleneck:
|
|
174
|
+
|
|
175
|
+
| Approach | Basis cost | Why |
|
|
176
|
+
|---|---|---|
|
|
177
|
+
| Cox-de Boor (PyKAN) | 3 sequential GPU passes | Each pass depends on previous |
|
|
178
|
+
| Gaussian RBF (FastKAN) | 1 `exp()` call | Fast but not a true B-spline |
|
|
179
|
+
| **Truncated power (FlashKAN)** | **1 fused kernel** | `clamp + multiply` is cheaper than `exp()` |
|
|
180
|
+
|
|
181
|
+
## B-spline properties preserved
|
|
182
|
+
|
|
183
|
+
Unlike Gaussian RBF approximations, FlashKAN computes **exact** B-spline basis values:
|
|
184
|
+
|
|
185
|
+
- **Compact support** — each basis function is exactly zero outside its knot span window
|
|
186
|
+
- **C2 continuity** — second derivatives are continuous at every knot
|
|
187
|
+
- **Partition of unity** — basis values sum to 1 at every point in the interior
|
|
188
|
+
- **Non-negativity** — all basis values are >= 0
|
|
189
|
+
|
|
190
|
+
Verified: max difference vs Cox-de Boor reference is < 5e-5 in float32.
|
|
191
|
+
|
|
192
|
+
## Project structure
|
|
193
|
+
|
|
194
|
+
```
|
|
195
|
+
src/flashkan/
|
|
196
|
+
├── __init__.py # Public API
|
|
197
|
+
├── basis.py # Truncated power B-spline + torch.compile (core math)
|
|
198
|
+
├── layer.py # KANLayer
|
|
199
|
+
├── network.py # KANNetwork
|
|
200
|
+
└── visualize.py # plot_basis, plot_activations, plot_network
|
|
201
|
+
```
|
|
202
|
+
|
|
203
|
+
5 source files. The core innovation is in `basis.py` — 30 lines of math.
|
|
204
|
+
|
|
205
|
+
## Citation
|
|
206
|
+
|
|
207
|
+
If you use FlashKAN in your research, please cite:
|
|
208
|
+
|
|
209
|
+
```bibtex
|
|
210
|
+
@software{flashkan2026,
|
|
211
|
+
title={FlashKAN: Fast B-spline KAN Layers via Truncated Power Basis},
|
|
212
|
+
author={Mysore, Naveen},
|
|
213
|
+
year={2026},
|
|
214
|
+
url={https://github.com/NAVEENMN/flashkan}
|
|
215
|
+
}
|
|
216
|
+
```
|
|
217
|
+
|
|
218
|
+
## License
|
|
219
|
+
|
|
220
|
+
MIT
|
inkan-0.1.0/README.md
ADDED
|
@@ -0,0 +1,192 @@
|
|
|
1
|
+
# FlashKAN
|
|
2
|
+
|
|
3
|
+
Fast B-spline [Kolmogorov-Arnold Network](https://arxiv.org/abs/2404.19756) layers for PyTorch.
|
|
4
|
+
|
|
5
|
+
**6-15x faster** than standard Cox-de Boor implementations (PyKAN, efficient-kan), **faster than Gaussian RBF** alternatives (FastKAN) — while producing **exact B-spline basis values** with compact support, C2 continuity, and partition of unity.
|
|
6
|
+
|
|
7
|
+
## How it works
|
|
8
|
+
|
|
9
|
+
Standard KAN implementations compute B-spline basis functions using the Cox-de Boor recursion — 3 sequential passes for cubic splines, each creating intermediate tensors. FlashKAN replaces this with the truncated power closed form:
|
|
10
|
+
|
|
11
|
+
```
|
|
12
|
+
N(u) = (1/6) [relu(u)³ - 4·relu(u-1)³ + 6·relu(u-2)³ - 4·relu(u-3)³ + relu(u-4)³]
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
This single expression computes exact B-spline values with no recursion, no span lookups, and no gather operations. `torch.compile` fuses all elementwise ops into one GPU kernel.
|
|
16
|
+
|
|
17
|
+
## Installation
|
|
18
|
+
|
|
19
|
+
```bash
|
|
20
|
+
pip install flashkan
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
**Requirements:** Python >= 3.9, PyTorch >= 2.0
|
|
24
|
+
|
|
25
|
+
**Supported devices:** CPU, CUDA (NVIDIA), MPS (Apple Silicon)
|
|
26
|
+
|
|
27
|
+
### From source
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
git clone https://github.com/NAVEENMN/flashkan.git
|
|
31
|
+
cd flashkan
|
|
32
|
+
pip install -e .
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
## Quick start
|
|
36
|
+
|
|
37
|
+
```python
|
|
38
|
+
import torch
|
|
39
|
+
from flashkan import KANLayer, KANNetwork
|
|
40
|
+
|
|
41
|
+
# Drop-in replacement for nn.Linear
|
|
42
|
+
layer = KANLayer(784, 64)
|
|
43
|
+
x = torch.randn(32, 784)
|
|
44
|
+
y = layer(x) # [32, 64]
|
|
45
|
+
|
|
46
|
+
# Multi-layer network
|
|
47
|
+
net = KANNetwork([784, 64, 10])
|
|
48
|
+
y = net(torch.randn(32, 784)) # [32, 10]
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
### MNIST example
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
import torch
|
|
55
|
+
import torch.nn as nn
|
|
56
|
+
from flashkan import KANNetwork
|
|
57
|
+
|
|
58
|
+
model = KANNetwork([784, 64, 10])
|
|
59
|
+
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
|
|
60
|
+
criterion = nn.CrossEntropyLoss()
|
|
61
|
+
|
|
62
|
+
# Standard PyTorch training loop
|
|
63
|
+
for images, labels in train_loader:
|
|
64
|
+
output = model(images.view(-1, 784))
|
|
65
|
+
loss = criterion(output, labels)
|
|
66
|
+
optimizer.zero_grad()
|
|
67
|
+
loss.backward()
|
|
68
|
+
optimizer.step()
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
See [`examples/`](examples/) for complete runnable scripts.
|
|
72
|
+
|
|
73
|
+
## Visualization
|
|
74
|
+
|
|
75
|
+
FlashKAN includes built-in visualization for learned activation functions — similar to PyKAN's `model.plot()`.
|
|
76
|
+
|
|
77
|
+
```python
|
|
78
|
+
from flashkan import KANNetwork, plot_basis, plot_activations, plot_network
|
|
79
|
+
|
|
80
|
+
model = KANNetwork([784, 32, 10], grid_size=8)
|
|
81
|
+
# ... train on MNIST ...
|
|
82
|
+
|
|
83
|
+
plot_basis(model.layers[0]) # B-spline basis bumps
|
|
84
|
+
plot_activations(model.layers[1]) # learned curves per edge
|
|
85
|
+
plot_network(model) # full network diagram
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
### B-spline basis functions
|
|
89
|
+
|
|
90
|
+
The 8 basis bumps (grid_size=5, degree=3) — compact support, smooth overlap:
|
|
91
|
+
|
|
92
|
+

|
|
93
|
+
|
|
94
|
+
### Learned activation functions
|
|
95
|
+
|
|
96
|
+
After training on MNIST, each edge learns a unique activation curve.
|
|
97
|
+
Cyan = total, red dashed = spline component, green dotted = SiLU base:
|
|
98
|
+
|
|
99
|
+

|
|
100
|
+
|
|
101
|
+
### Network diagram
|
|
102
|
+
|
|
103
|
+
Full [784 → 32 → 10] network with learned curves on edges:
|
|
104
|
+
|
|
105
|
+

|
|
106
|
+
|
|
107
|
+
## API
|
|
108
|
+
|
|
109
|
+
### `KANLayer(in_features, out_features, grid_size=5, spline_order=3)`
|
|
110
|
+
|
|
111
|
+
A single KAN layer. Drop-in replacement for `nn.Linear`.
|
|
112
|
+
|
|
113
|
+
| Parameter | Default | Description |
|
|
114
|
+
|---|---|---|
|
|
115
|
+
| `in_features` | — | Input dimension |
|
|
116
|
+
| `out_features` | — | Output dimension |
|
|
117
|
+
| `grid_size` | 5 | Number of knot intervals (more = finer approximation) |
|
|
118
|
+
| `spline_order` | 3 | B-spline degree (3 = cubic, recommended) |
|
|
119
|
+
| `grid_range` | (-1, 1) | Input range for the spline grid |
|
|
120
|
+
|
|
121
|
+
### `KANNetwork(layer_dims, grid_size=5, spline_order=3)`
|
|
122
|
+
|
|
123
|
+
Stack of KAN layers.
|
|
124
|
+
|
|
125
|
+
```python
|
|
126
|
+
# 3-layer KAN: 784 -> 128 -> 64 -> 10
|
|
127
|
+
net = KANNetwork([784, 128, 64, 10])
|
|
128
|
+
```
|
|
129
|
+
|
|
130
|
+
## Benchmarks
|
|
131
|
+
|
|
132
|
+
Forward pass time (ms) on Apple M-series GPU (MPS), batch=256:
|
|
133
|
+
|
|
134
|
+
| Layer | MNIST (784→64) | FashionMNIST (784→64) | CIFAR-10 (3072→64) |
|
|
135
|
+
|---|---|---|---|
|
|
136
|
+
| **FlashKAN (compiled)** | **0.27** | **0.20** | **0.39** |
|
|
137
|
+
| FastKAN (Gaussian RBF) | 0.38 | 0.31 | 0.99 |
|
|
138
|
+
| NoGather (unrolled) | 0.99 | 1.02 | 4.85 |
|
|
139
|
+
| Vanilla (Cox-de Boor) | 1.69 | 1.67 | 5.98 |
|
|
140
|
+
|
|
141
|
+
FlashKAN is **6.3x faster** than vanilla Cox-de Boor on MNIST and **15.3x faster** on CIFAR-10.
|
|
142
|
+
|
|
143
|
+
### Why it's fast
|
|
144
|
+
|
|
145
|
+
91% of a standard KAN forward pass is spent computing B-spline basis functions. FlashKAN eliminates this bottleneck:
|
|
146
|
+
|
|
147
|
+
| Approach | Basis cost | Why |
|
|
148
|
+
|---|---|---|
|
|
149
|
+
| Cox-de Boor (PyKAN) | 3 sequential GPU passes | Each pass depends on previous |
|
|
150
|
+
| Gaussian RBF (FastKAN) | 1 `exp()` call | Fast but not a true B-spline |
|
|
151
|
+
| **Truncated power (FlashKAN)** | **1 fused kernel** | `clamp + multiply` is cheaper than `exp()` |
|
|
152
|
+
|
|
153
|
+
## B-spline properties preserved
|
|
154
|
+
|
|
155
|
+
Unlike Gaussian RBF approximations, FlashKAN computes **exact** B-spline basis values:
|
|
156
|
+
|
|
157
|
+
- **Compact support** — each basis function is exactly zero outside its knot span window
|
|
158
|
+
- **C2 continuity** — second derivatives are continuous at every knot
|
|
159
|
+
- **Partition of unity** — basis values sum to 1 at every point in the interior
|
|
160
|
+
- **Non-negativity** — all basis values are >= 0
|
|
161
|
+
|
|
162
|
+
Verified: max difference vs Cox-de Boor reference is < 5e-5 in float32.
|
|
163
|
+
|
|
164
|
+
## Project structure
|
|
165
|
+
|
|
166
|
+
```
|
|
167
|
+
src/flashkan/
|
|
168
|
+
├── __init__.py # Public API
|
|
169
|
+
├── basis.py # Truncated power B-spline + torch.compile (core math)
|
|
170
|
+
├── layer.py # KANLayer
|
|
171
|
+
├── network.py # KANNetwork
|
|
172
|
+
└── visualize.py # plot_basis, plot_activations, plot_network
|
|
173
|
+
```
|
|
174
|
+
|
|
175
|
+
5 source files. The core innovation is in `basis.py` — 30 lines of math.
|
|
176
|
+
|
|
177
|
+
## Citation
|
|
178
|
+
|
|
179
|
+
If you use FlashKAN in your research, please cite:
|
|
180
|
+
|
|
181
|
+
```bibtex
|
|
182
|
+
@software{flashkan2026,
|
|
183
|
+
title={FlashKAN: Fast B-spline KAN Layers via Truncated Power Basis},
|
|
184
|
+
author={Mysore, Naveen},
|
|
185
|
+
year={2026},
|
|
186
|
+
url={https://github.com/NAVEENMN/flashkan}
|
|
187
|
+
}
|
|
188
|
+
```
|
|
189
|
+
|
|
190
|
+
## License
|
|
191
|
+
|
|
192
|
+
MIT
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68.0", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "inkan"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Fast B-spline KAN layers — 6-15x faster than Cox-de Boor, exact B-spline values"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = {text = "MIT"}
|
|
11
|
+
requires-python = ">=3.9"
|
|
12
|
+
authors = [
|
|
13
|
+
{name = "Naveen Mysore"},
|
|
14
|
+
]
|
|
15
|
+
keywords = ["kan", "kolmogorov-arnold", "b-spline", "neural-network", "deep-learning"]
|
|
16
|
+
classifiers = [
|
|
17
|
+
"Development Status :: 3 - Alpha",
|
|
18
|
+
"Intended Audience :: Science/Research",
|
|
19
|
+
"License :: OSI Approved :: MIT License",
|
|
20
|
+
"Programming Language :: Python :: 3",
|
|
21
|
+
"Programming Language :: Python :: 3.9",
|
|
22
|
+
"Programming Language :: Python :: 3.10",
|
|
23
|
+
"Programming Language :: Python :: 3.11",
|
|
24
|
+
"Programming Language :: Python :: 3.12",
|
|
25
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
26
|
+
]
|
|
27
|
+
dependencies = [
|
|
28
|
+
"torch>=2.0",
|
|
29
|
+
]
|
|
30
|
+
|
|
31
|
+
[project.optional-dependencies]
|
|
32
|
+
dev = [
|
|
33
|
+
"pytest>=7.0",
|
|
34
|
+
"torchvision>=0.15",
|
|
35
|
+
"matplotlib>=3.5",
|
|
36
|
+
]
|
|
37
|
+
|
|
38
|
+
[project.urls]
|
|
39
|
+
Homepage = "https://github.com/naveenmysore/flashkan"
|
|
40
|
+
Repository = "https://github.com/naveenmysore/flashkan"
|
|
41
|
+
|
|
42
|
+
[tool.setuptools.packages.find]
|
|
43
|
+
where = ["src"]
|
inkan-0.1.0/setup.cfg
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
"""InKAN — Fast B-spline KAN layers using the truncated power basis.
|
|
2
|
+
|
|
3
|
+
Computes exact B-spline basis functions via a closed-form formula
|
|
4
|
+
that torch.compile fuses into a single GPU kernel. 6-15x faster
|
|
5
|
+
than Cox-de Boor recursion, faster than Gaussian RBF alternatives,
|
|
6
|
+
while maintaining true B-spline properties (compact support, C2
|
|
7
|
+
continuity, partition of unity).
|
|
8
|
+
|
|
9
|
+
Example:
|
|
10
|
+
>>> import torch
|
|
11
|
+
>>> from inkan import KANLayer
|
|
12
|
+
>>> layer = KANLayer(784, 64)
|
|
13
|
+
>>> x = torch.randn(32, 784)
|
|
14
|
+
>>> y = layer(x) # [32, 64]
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
__version__ = "0.2.0"
|
|
18
|
+
|
|
19
|
+
from inkan.layer import KANLayer
|
|
20
|
+
from inkan.network import KANNetwork
|
|
21
|
+
from inkan.visualize import plot_basis, plot_activations, plot_network
|
|
22
|
+
|
|
23
|
+
__all__ = ["KANLayer", "KANNetwork",
|
|
24
|
+
"plot_basis", "plot_activations", "plot_network"]
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
"""B-spline basis computation via the truncated power closed form.
|
|
2
|
+
|
|
3
|
+
The cubic B-spline basis function can be expressed exactly as:
|
|
4
|
+
|
|
5
|
+
N(u) = (1/6) * [relu(u)^3 - 4*relu(u-1)^3 + 6*relu(u-2)^3
|
|
6
|
+
- 4*relu(u-3)^3 + relu(u-4)^3]
|
|
7
|
+
|
|
8
|
+
where u = (x - grid_start_i) / h.
|
|
9
|
+
|
|
10
|
+
This eliminates Cox-de Boor recursion (3 sequential passes),
|
|
11
|
+
span lookups, and scatter/gather operations. torch.compile fuses
|
|
12
|
+
all elementwise ops into a single GPU kernel.
|
|
13
|
+
|
|
14
|
+
For uniform grids, this produces bit-identical results to Cox-de Boor
|
|
15
|
+
(verified max diff < 5e-5 in float32).
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
import torch
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def _bspline_basis(x: torch.Tensor, grid_starts: torch.Tensor,
|
|
22
|
+
inv_h: float) -> torch.Tensor:
|
|
23
|
+
"""Compute B-spline basis values for all inputs and all basis functions.
|
|
24
|
+
|
|
25
|
+
Args:
|
|
26
|
+
x: Input tensor [batch, in_features].
|
|
27
|
+
grid_starts: Start position of each basis function's support [n_bases].
|
|
28
|
+
inv_h: Reciprocal of knot spacing (1/h).
|
|
29
|
+
|
|
30
|
+
Returns:
|
|
31
|
+
Basis values [batch, in_features, n_bases].
|
|
32
|
+
"""
|
|
33
|
+
# u[b, i, j] = (x[b,i] - grid_starts[j]) / h
|
|
34
|
+
u = (x.unsqueeze(-1) - grid_starts) * inv_h
|
|
35
|
+
|
|
36
|
+
# Truncated power form: sum of relu(u-k)^3 with binomial coefficients
|
|
37
|
+
# Coefficients: C(4,k) * (-1)^k / 3! = [1, -4, 6, -4, 1] / 6
|
|
38
|
+
r0 = torch.clamp(u, min=0.0)
|
|
39
|
+
r1 = torch.clamp(u - 1.0, min=0.0)
|
|
40
|
+
r2 = torch.clamp(u - 2.0, min=0.0)
|
|
41
|
+
r3 = torch.clamp(u - 3.0, min=0.0)
|
|
42
|
+
r4 = torch.clamp(u - 4.0, min=0.0)
|
|
43
|
+
|
|
44
|
+
return (r0*r0*r0 - 4.0*r1*r1*r1 + 6.0*r2*r2*r2
|
|
45
|
+
- 4.0*r3*r3*r3 + r4*r4*r4) * 0.16666667
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
# torch.compile fuses all elementwise ops into a single GPU kernel.
|
|
49
|
+
# Works on CUDA, MPS, and CPU.
|
|
50
|
+
bspline_basis = torch.compile(_bspline_basis)
|
|
51
|
+
|
|
52
|
+
# Export the unfused version for testing/debugging
|
|
53
|
+
bspline_basis_eager = _bspline_basis
|
|
@@ -0,0 +1,112 @@
|
|
|
1
|
+
"""KANLayer — the core building block.
|
|
2
|
+
|
|
3
|
+
Drop-in replacement for nn.Linear in KAN networks:
|
|
4
|
+
nn.Linear(in_features, out_features)
|
|
5
|
+
KANLayer(in_features, out_features)
|
|
6
|
+
|
|
7
|
+
Each edge (i→j) has a learnable B-spline activation function
|
|
8
|
+
plus a SiLU residual connection, following the KAN paper.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
import torch
|
|
12
|
+
import torch.nn as nn
|
|
13
|
+
import torch.nn.functional as F
|
|
14
|
+
|
|
15
|
+
from inkan.basis import bspline_basis
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class KANLayer(nn.Module):
|
|
19
|
+
"""Kolmogorov-Arnold Network layer with fast B-spline activations.
|
|
20
|
+
|
|
21
|
+
Uses the truncated power closed form + torch.compile for basis
|
|
22
|
+
computation. Produces exact B-spline values with compact support,
|
|
23
|
+
C2 continuity, and partition of unity.
|
|
24
|
+
|
|
25
|
+
Args:
|
|
26
|
+
in_features: Size of each input sample.
|
|
27
|
+
out_features: Size of each output sample.
|
|
28
|
+
grid_size: Number of interior knot intervals. More = finer
|
|
29
|
+
approximation. Default: 5.
|
|
30
|
+
spline_order: Degree of the B-spline. Default: 3 (cubic).
|
|
31
|
+
base_activation: Residual activation function. Default: SiLU.
|
|
32
|
+
grid_range: Range of the input grid. Default: (-1, 1).
|
|
33
|
+
|
|
34
|
+
Shape:
|
|
35
|
+
- Input: (batch, in_features)
|
|
36
|
+
- Output: (batch, out_features)
|
|
37
|
+
|
|
38
|
+
Example:
|
|
39
|
+
>>> layer = KANLayer(784, 64, grid_size=10)
|
|
40
|
+
>>> x = torch.randn(32, 784)
|
|
41
|
+
>>> y = layer(x) # [32, 64]
|
|
42
|
+
"""
|
|
43
|
+
|
|
44
|
+
def __init__(
|
|
45
|
+
self,
|
|
46
|
+
in_features: int,
|
|
47
|
+
out_features: int,
|
|
48
|
+
grid_size: int = 5,
|
|
49
|
+
spline_order: int = 3,
|
|
50
|
+
base_activation: type = nn.SiLU,
|
|
51
|
+
grid_range: tuple = (-1.0, 1.0),
|
|
52
|
+
):
|
|
53
|
+
super().__init__()
|
|
54
|
+
self.in_features = in_features
|
|
55
|
+
self.out_features = out_features
|
|
56
|
+
self.grid_size = grid_size
|
|
57
|
+
self.spline_order = spline_order
|
|
58
|
+
|
|
59
|
+
n_bases = grid_size + spline_order
|
|
60
|
+
h = (grid_range[1] - grid_range[0]) / grid_size
|
|
61
|
+
self.inv_h = 1.0 / h
|
|
62
|
+
|
|
63
|
+
# Grid start positions for each basis function
|
|
64
|
+
grid_starts = (torch.arange(n_bases).float() * h
|
|
65
|
+
+ grid_range[0] - spline_order * h)
|
|
66
|
+
self.register_buffer("grid_starts", grid_starts)
|
|
67
|
+
|
|
68
|
+
# Learnable spline coefficients: one per (output, input, basis)
|
|
69
|
+
self.spline_weight = nn.Parameter(
|
|
70
|
+
torch.empty(out_features, in_features, n_bases)
|
|
71
|
+
)
|
|
72
|
+
|
|
73
|
+
# Learnable residual weights
|
|
74
|
+
self.base_weight = nn.Parameter(
|
|
75
|
+
torch.empty(out_features, in_features)
|
|
76
|
+
)
|
|
77
|
+
|
|
78
|
+
self.base_activation = base_activation()
|
|
79
|
+
|
|
80
|
+
self.reset_parameters()
|
|
81
|
+
|
|
82
|
+
def reset_parameters(self):
|
|
83
|
+
"""Initialize parameters with scaled random values."""
|
|
84
|
+
# Kaiming-style init scaled for spline basis
|
|
85
|
+
nn.init.trunc_normal_(self.spline_weight, std=0.1)
|
|
86
|
+
nn.init.trunc_normal_(self.base_weight, std=0.1)
|
|
87
|
+
|
|
88
|
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
89
|
+
"""Forward pass.
|
|
90
|
+
|
|
91
|
+
Args:
|
|
92
|
+
x: Input tensor of shape (batch, in_features).
|
|
93
|
+
|
|
94
|
+
Returns:
|
|
95
|
+
Output tensor of shape (batch, out_features).
|
|
96
|
+
"""
|
|
97
|
+
# Spline activation: φ(x) = Σ_k c_k * N_k(x)
|
|
98
|
+
bases = bspline_basis(x, self.grid_starts, self.inv_h)
|
|
99
|
+
spline_out = torch.einsum("bin,oin->bo", bases, self.spline_weight)
|
|
100
|
+
|
|
101
|
+
# Residual: b(x) * w
|
|
102
|
+
base_out = torch.einsum(
|
|
103
|
+
"bi,oi->bo", self.base_activation(x), self.base_weight
|
|
104
|
+
)
|
|
105
|
+
|
|
106
|
+
return spline_out + base_out
|
|
107
|
+
|
|
108
|
+
def extra_repr(self) -> str:
|
|
109
|
+
return (f"in_features={self.in_features}, "
|
|
110
|
+
f"out_features={self.out_features}, "
|
|
111
|
+
f"grid_size={self.grid_size}, "
|
|
112
|
+
f"spline_order={self.spline_order}")
|