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 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
+ ![Basis functions](assets/basis.png)
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
+ ![Learned activations](assets/activations.png)
128
+
129
+ ### Network diagram
130
+
131
+ Full [784 → 32 → 10] network with learned curves on edges:
132
+
133
+ ![Network diagram](assets/network.png)
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
+ ![Basis functions](assets/basis.png)
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
+ ![Learned activations](assets/activations.png)
100
+
101
+ ### Network diagram
102
+
103
+ Full [784 → 32 → 10] network with learned curves on edges:
104
+
105
+ ![Network diagram](assets/network.png)
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,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -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}")