JAX¶
kanforge.jax.layers
¶
JAX KAN layers on Flax NNX. Run on CPU, CUDA/ROCm GPUs and TPUs.
Same parameter names, shapes and math as kanforge.torch and
kanforge.mlx. Trainable arrays are nnx.Param; grids, masks and
symbolic edges are Buffer variables, which nnx.grad and optimizers
(wrt=nnx.Param) leave alone. Fused Pallas kernels plug in behind
backend (see kanforge.jax.kernels).
Buffer
¶
Bases: Variable
Non-trainable layer state (grids, edge masks, symbolic edges).
KANLinear(in_features, out_features, basis='bspline', *, rngs=None, bias=True, base_activation=jax.nn.silu, groups=None, edge_scale=False, grid_ranges=None, coef_init='normal', coef_init_std=0.1, edge_dropout=0.0, backend='auto', dtype=jnp.float32)
¶
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.
rngs initialises the parameters and drives edge_dropout
(stream dropout); model.eval() switches dropout off.
uses_kernel()
¶
Whether forward runs the fused Pallas kernel (on the default JAX backend).
edge_derivative(x, order=1)
¶
d^order phi_oi / dx_i^order for every edge, shape (..., out, in).
state_dict()
¶
NumPy copies of every variable, keyed like torch.nn.Module.state_dict.
from_linear(linear, activation=jax.nn.silu, basis='bspline', **kwargs)
classmethod
¶
Exactly linear(activation(x)) with zero spline coefficients.
GRKAN(in_features, out_features, *, groups=8, num_degree=5, den_degree=4, init='gelu', bias=True, rngs=None, dtype=jnp.float32)
¶
KAN(widths, basis='bspline', grid_ranges=None, *, rngs=None, layers=None, **layer_kwargs)
¶
Bases: Module
Stack of KANLinear layers; grid_ranges sets the first layer's grids.
MultKANLinear(in_features, n_sum, n_mult=0, basis='bspline', *, arity=2, rngs=None, **kwargs)
¶
Bases: Module
KAN layer with addition and multiplication nodes (see kanforge.torch).
from_mlp(mlp, basis='bspline', **kwargs)
¶
JAX version of kanforge.torch.from_mlp (layers of an nnx.Sequential).
rebuild_kan(layers)
¶
A KAN made of existing layers (used by pruning).
layer_inputs(model, x)
¶
(KANLinear, its input) for KANLinear, MultKANLinear, KAN, MultKAN, nnx.Sequential.
regularization(model, x=None, *, l1=1.0, entropy=2.0, coef_l1=0.0, coef_smooth=0.0)
¶
pykan's sparsity penalty (see kanforge.torch.regularization).
kanforge.jax.kernels
¶
Pallas kernels for KAN layers: the F / G / P primitives of kanforge.kernels.
Each primitive is a jax.custom_vjp whose backward pass is written with
the other primitives, so jax.grad nests to any order (PINN residuals).
Plans, folding of polynomial / Fourier derivatives and the feature formulas
are those of the Metal and Triton kernels. Shapes are padded to powers of two
(Pallas' Triton lowering needs them); padding carries zero weights.
Runs on GPU (Pallas -> Triton) and, with KANFORGE_PALLAS_INTERPRET=1,
on any backend in Pallas' interpreter (used by the tests).
spline_forward(x, coef, grid, plan)
¶
sum_{i,k} coef[o,i,k] phi_k(x[..., i]) through the Pallas kernel.