PyTorch¶
kanforge.torch.layers
¶
PyTorch KAN layers. Run on any torch device: CUDA, ROCm, MPS, CPU.
This is the plain-ops implementation; fused kernels will plug in behind the same modules without changing parameters or state_dict layout.
KANLinear(in_features, out_features, basis='bspline', *, bias=True, base_activation=F.silu, groups=None, edge_scale=False, grid_ranges=None, coef_init='normal', coef_init_std=0.1, edge_dropout=0.0, backend='auto', device=None, dtype=None)
¶
Bases: GridToolsMixin, InterpretMixin, Module
One KAN layer: y_o = sum_i s_oi * phi_oi(x_i) + sum_i w_oi * act(x_i) + b_o.
Parameters (same names on every backend):
coef (out, in, basis.size) — basis coefficients of every edge
base_weight (out, in) — residual branch, absent if base_activation is None
bias (out,) — absent if bias=False
edge_scale (out, in) — per-edge scale s_oi, only with edge_scale=True
basis_<name> trainable basis parameters (groups, *shape): groups
inputs share them (None: one set per input, 1: one per layer)
Buffers:
basis_<name> basis state such as grids, (in, *shape): every input
has its own grid (grid_ranges, update_grid)
edge_mask, sym_code, sym_params — pruning and symbolic edges
(see InterpretMixin, kanforge.interpret)
edge_dropout drops individual edge activations while training (DropKAN).
backend: "ops" — plain PyTorch ops; "kernel" — fused Metal (MPS)
or Triton (CUDA/ROCm) kernel, an error if the basis or device has none;
"auto" — the kernel where benchmarks showed it faster (AUTO_KINDS).
Both give the same values and derivatives of every order.
basis_params()
¶
Basis parameters and buffers expanded to (in, *shape) (or (1, *shape)).
uses_kernel(device)
¶
Whether forward runs the fused kernel on device.
edge_derivative(x, order=1)
¶
d^order phi_oi / dx_i^order for every edge, shape (..., out, in).
order=0 gives the edge functions themselves (bias excluded). A KAN
layer is a sum of 1-D edge functions, so entry (o, i) is exactly
d^m y_o / dx_i^m — the diagonal derivatives PINN losses need.
from_linear(linear, activation=F.silu, basis='bspline', **kwargs)
classmethod
¶
Exactly linear(activation(x)): the activation becomes the base branch,
spline coefficients start at zero and learn corrections.
GRKAN(in_features, out_features, *, groups=8, num_degree=5, den_degree=4, init='gelu', bias=True, device=None, dtype=None)
¶
Bases: KANLinear
Group-rational KAN layer of KAT: rational activations shared by groups
of inputs, then the layer coefficients act as the linear weight; no base branch.
KAN(widths, basis='bspline', grid_ranges=None, **layer_kwargs)
¶
Bases: Sequential
Stack of KANLinear layers, e.g. KAN([2, 5, 1], basis="chebyshev").
grid_ranges (one (lo, hi) per model input) sets the grids of the first
layer; hidden layers keep the basis default (adapt them with update_grid).
MultKANLinear(in_features, n_sum, n_mult=0, basis='bspline', *, arity=2, **kwargs)
¶
Bases: Module
KAN layer with n_sum addition nodes and n_mult multiplication nodes
(MultKAN, KAN 2.0). Every product node multiplies arity sub-nodes.
Output: (..., n_sum + n_mult).
MultKAN(widths, basis='bspline', *, arity=2, **layer_kwargs)
¶
Bases: Sequential
MultKAN([2, (2, 1), 1]): an int is a width of addition nodes,
(n_sum, n_mult) adds multiplication nodes to that layer.
from_mlp(mlp, basis='bspline', **kwargs)
¶
Turn Linear, act, Linear, act, ..., Linear into an equal network:
the first Linear stays, every (act, Linear) pair becomes a KANLinear
whose base branch is that activation. Outputs match exactly at start.
rebuild_kan(layers)
¶
A KAN made of existing layers (used by pruning).
layer_inputs(model, x)
¶
(KANLinear, the input it receives) for every KAN layer of model.
Works for KANLinear, MultKANLinear and any nn.Sequential of modules.
regularization(model, x=None, *, l1=1.0, entropy=2.0, coef_l1=0.0, coef_smooth=0.0)
¶
pykan's sparsity penalty summed over all KAN layers.
With x the edge magnitudes are mean |phi_oi| over the batch (pykan);
without it, mean |coef| (efficient-kan). Add lamb * regularization(...)
to the loss.
kanforge.torch.conv
¶
KAN convolutions (channels-first, like torch.nn.ConvNd).
Every output pixel is a KANLinear of its receptive field: each element of
the kernel is a learnable edge function instead of a weight (as in
torch-conv-kan and Convolutional-KANs). Patch features are ordered like the
weights of nn.ConvNd — (in_channels/groups, k_1, ..., k_n) — so an
ordinary convolution converts exactly with from_conv.
kanforge.torch.transformer
¶
KAN feed-forward blocks for transformers and model surgery for Hugging Face.
KANFFN— the MLP of a transformer block made of KAN layers: either twoKANLinearwith any basis, or KAT's two group-rational layers (the first rational starts as the identity, the second as the activation).KATBlock— pre-norm transformer block: attention +KANFFN.kanify(model)— replace the FFNs of a pretrained model in place:mode="exact"keeps the model's outputs bit for bit at the start (the activation becomes the base branch of aKANLinear, splines start at zero),mode="grkan"transfers the weights into GR-KAN layers like KAT does from ViT (rationals approximate the activation to ~1e-2).
KANFFN(dim, hidden=None, kind='grkan', basis='bspline', groups=8, activation='gelu', dropout=0.0, **kan_kwargs)
¶
Bases: Module
dim -> hidden -> dim with KAN layers (kind="kan") or GR-KAN ("grkan").
KATBlock(dim, heads, mlp_ratio=4.0, kind='grkan', dropout=0.0, **ffn_kwargs)
¶
Bases: Module
x + attn(norm(x)); x + KANFFN(norm(x)) — the block of the Kolmogorov–Arnold Transformer.
kanify(model, mode='exact', basis='bspline', groups=8, **kan_kwargs)
¶
Replace the feed-forward blocks of model in place; returns their names.
mode="exact": the second projection becomes KANLinear.from_linear
with the activation as base branch and the activation slot an identity;
gated FFNs get a KANLinear down projection with an identity base. The
model computes exactly the same function until the splines are trained.
mode="grkan": fc1 -> GR-KAN with identity rationals, fc2 -> GR-KAN with
rationals fitted to the activation (KAT's weight transfer); approximate.
Gated FFNs are left alone in this mode.
kanforge.torch.graph
¶
Graph KAN layers (KAGNNs): message passing with KAN transformations.
Graphs are given as edge_index (2, E) with rows (source, target), the
PyTorch Geometric convention, so the layers drop into PyG pipelines but do
not need PyG.
KAGINConv— h_v' = KAN((1 + eps) h_v + sum_{u -> v} h_u) (GIN with a KAN).KAGCNConv— h' = KAN(D^-1/2 (A + I) D^-1/2 h) (GCN aggregation, then KAN).KAGATConv— attention over neighbours of KAN-transformed features (GAT).
KAGATConv(in_features, out_features, heads=1, basis='bspline', concat=True, negative_slope=0.2, add_loops=True, **kan_kwargs)
¶
Bases: Module
Graph attention with a KAN feature transform; heads concatenated or averaged.
aggregate(x, edge_index, num_nodes, weight=None)
¶
out[v] = sum over edges (u -> v) of weight * x[u].
kanforge.torch.recurrent
¶
Recurrent KAN cells: LSTM and GRU whose gates are KAN layers (in the spirit of TKAN).
The gate pre-activations come from a KANLinear of [x_t, h_{t-1}] instead
of an affine map; the gating itself is the usual LSTM / GRU algebra.
KANRNN runs a stack of cells over a sequence.
kanforge.torch.science
¶
Scientific machine learning with KANs: PINN tools, operators, ODEs.
derivative/gradient/laplacian— derivatives of a field u(x) for PDE residuals, kept in the graph for training (any model; KAN layers additionally have exactedge_derivative).sample_box/sample_boundary— collocation points.DeepONet— operator learning with KAN branch and trunk nets (DeepOKAN).FBKAN— finite-basis KAN: overlapping subdomains with a partition of unity, one KAN per subdomain with its grid on that subdomain.KANODE— neural ODE dy/dt = KAN(y) (optionally with t) integrated by fixed-step RK4, differentiable end to end (KAN-ODEs).
DeepONet(branch_widths, trunk_widths, basis='bspline', **kan_kwargs)
¶
Bases: Module
G(u)(y) = sum_k branch_k(u) * trunk_k(y) + b with KAN branch and trunk.
forward(u, y)
¶
u (B, m) sensor values, y (M, d) query points -> (B, M).
FBKAN(widths, bounds, parts=2, overlap=0.25, basis='bspline', **kan_kwargs)
¶
Bases: Module
Finite-basis KAN (domain decomposition, as in NeuroMANCER's FBKAN).
The box is split into parts subdomains per dimension, each enlarged by
overlap (fraction of its width). Window w_j(x) = prod_d cos^2 bump on
the enlarged subdomain; u(x) = sum_j w_j(x) KAN_j(x) / sum_j w_j(x).
windows(x)
¶
(N, J) normalised partition of unity.
KANODE(dim, hidden=(8,), time_dependent=False, basis='bspline', **kan_kwargs)
¶
Bases: Module
dy/dt = f(y) or f(t, y) with f a KAN, integrated by RK4.
forward(y0, t, substeps=1)
¶
Trajectory (T, B, dim) at times t (T,) starting from y0 at t[0].
derivative(fn, x, order=1, dim=0)
¶
d^order u / dx_dim^order for u = fn(x), x (N, d), u (N,) or (N, 1).
gradient(fn, x)
¶
(N, d) gradient of u = fn(x).
laplacian(fn, x)
¶
(N,) sum of second derivatives of u = fn(x).
sample_box(bounds, n, generator=None, **kw)
¶
n uniform points inside the box bounds [(lo, hi), ...].
sample_boundary(bounds, n, generator=None, **kw)
¶
n uniform points on the faces of the box (each face equally likely).
kanforge.torch.kernels
¶
PyTorch side of the fused kernels: launchers and differentiable primitives.
MPS tensors run the Metal kernels, CUDA/ROCm tensors the Triton kernels;
with TRITON_INTERPRET=1 the Triton kernels also run on CPU tensors
(used by the tests). See kanforge.kernels for the math.
spline_forward(x, coef, grid, plan)
¶
sum_{i,k} coef[o,i,k] phi_k(x[..., i]) through the fused kernel.