Export and conversion

kanforge.convert

Move models between PyTorch, MLX and JAX, and save / load them anywhere.

A model is described by a JSON-able spec (architecture, basis configs, constructor options, activation names) plus a flat dict of NumPy arrays named like torch's state_dict. Supported: KANLinear, GRKAN, KAN, MultKANLinear, MultKAN and sequences of those with Linear layers and standard activations (models from from_mlp).

mlx_model = convert(torch_model, "mlx")
save(model, "model.npz");  model = load("model.npz", backend="jax")

activation_name(act)

Portable name of a base activation / activation layer (None = absent).

make_activation(name, backend)

Callable for name on backend ("torch", "mlx", "jax").

to_spec(model)

(spec, arrays) of a model of any backend.

from_spec(spec, arrays, backend)

Build the model described by spec on backend and load arrays.

convert(model, backend)

The same model (architecture and weights) on another backend.

save(model, path)

Backend-independent file: arrays and the JSON spec in one .npz.

load(path, backend='torch')

Load a file written by save on any backend.

kanforge.export.lut

Lookup-table form of KAN models: the deployable representation.

Every edge phi_oi of a layer is tabulated at size uniform points over its input range; a layer then computes

t = (x_i - lo_i) / step_i,  j = clip(floor(t), 0, size - 2),  f = t - j
y_o = b_o + sum_i T[i, j, o] + f * (T[i, j+1, o] - T[i, j, o])

i.e. gather + lerp + sum — operations every runtime (ONNX, Core ML, FPGA logic) has. Outside the range every edge continues linearly with its end slope (extrapolate="linear", which follows the SiLU base branch; the slopes are float and estimated over the last 1/32 of the table, so quantisation noise does not leak into them) or is held ("clamp"). A table is only as good as its range: choose margin / samples so that inputs stay inside. Works for models of any backend; MultKAN product nodes, Linear layers and standard activations are kept as such.

quantize(lut, bits) stores the tables as signed integers with one scale per output and a per-edge offset folded into the bias, so a layer sums integers and multiplies once per output.

LUTLayer(tables, lo, hi, bias, n_sum=None, n_mult=0, arity=2, scale=None, bits=None, slope_lo=None, slope_hi=None) dataclass

Tables (in, size, out) over per-input ranges [lo, hi].

end_slopes(tables)

(slope_lo, slope_hi) per index step, each (in, out), from float tables.

to_lut(model, x=None, size=256, margin=0.05, extrapolate='linear')

Tabulate a model (any backend). Ranges come from the activations every layer receives for samples x (widened by margin), or from the grids.

quantize(lut, bits=8)

Integer tables: T[i, :, o] = c_io + s_o * q[i, :, o], |q| <= 2^(bits-1) - 1.

kanforge.export.onnx_export

ONNX export.

• export_onnx(model_or_lut, path, x) — builds the graph directly from the lookup-table form (gather + lerp + sum per layer, int8 tables with a scale per output, product nodes, Linear layers, activations). Independent of the training backend: PyTorch, MLX and JAX models export the same way. Runs in onnxruntime and any ONNX runtime with opset 17.
• export_onnx_exact(torch_model, path, example) — the exact network (basis functions as arithmetic) through torch.onnx.export.

onnx_model(lut)

ONNX ModelProto of a lookup-table model; input "x" (batch, n_in) float32.

export_onnx(model, path, x=None, size=256, bits=None, margin=0.05, extrapolate='linear')

Write model (any backend, or a LUTModel) as an ONNX lookup-table network.

export_onnx_exact(model, path, example, opset=OPSET)

The exact PyTorch network as ONNX (torch.onnx.export, TorchScript exporter).

kanforge.export.coreml

Core ML export (coremltools): Apple devices, Neural Engine / GPU / CPU.

• export_coreml(model_or_lut, path, x) — the lookup-table network as an ML Program built op by op (gather + lerp + sum; int8 tables stored as int8 with constexpr_affine_dequantize), for models of any backend. Needs macOS 13 / iOS 16 or newer.
• export_coreml_exact(torch_model, path, example) — the exact network through torch.jit.trace and coremltools.convert.

coreml_model(lut, max_batch=65536)

coremltools.models.MLModel of a lookup-table model (input "x", output "y").

export_coreml(model, path, x=None, size=256, bits=None, margin=0.05, extrapolate='linear')

Write model (any backend, or a LUTModel) as a Core ML lookup-table network.

export_coreml_exact(model, path, example, max_batch=65536)

The exact PyTorch network as Core ML (traced).

kanforge.export.fpga

FPGA generator: KAN layers as integer lookup tables in Verilog (KANELÉ-style).

Every edge becomes a ROM indexed by its quantised input, so a layer is memory + adders, no multipliers in the datapath:

• inputs of a layer are A-bit addresses: a_i = clamp(round((x_i - lo_i) / step_i))
• edge (i, o) stores q[i, a, o] = round((phi_oi(lo_i + a step_i) - c_io) / s_o), W bits
• acc_o = sum_i q[i, a_i, o] (adder per output)
• real output y_o = s_o * acc_o + b_o (b_o includes sum_i c_io)
• next layer's address: a'_o = clamp((acc_o * M_o + K_o) >>> SH, 0, 2^A - 1) with integers M_o, K_o chosen so that a'_o = round((y_o - lo'_o) / step'_o)

One register stage per layer; the last layer outputs the accumulators and design.json holds s_o and b_o to turn them into real values. simulate(design, x) is the bit-exact NumPy model of the generated logic; the testbench compares the RTL against vectors from it.

design(model, x=None, address_bits=8, value_bits=8, margin=0.05)

Integer design of a KAN (any backend, or a float LUTModel with 2^A points).

quantize_input(d, x)

A-bit addresses of real inputs (what the host feeds the FPGA).

simulate(d, x=None, addresses=None)

Bit-exact model of the RTL: (last-layer accumulators, real outputs).

write_verilog(d, out_dir, test_inputs=None)

Write RTL, ROM images, a self-checking testbench and design.json.

export_fpga(model, out_dir, x=None, address_bits=8, value_bits=8, margin=0.05, test_vectors=64)

design + write_verilog with a testbench on test_vectors samples of x.

kanforge.grid

Grid tools, the same for every basis and backend (NumPy float64).

• adaptive_grid — pykan's data-driven grid: quantiles of the inputs mixed with a uniform grid over their range (grid_eps of uniform).
• refit — move a layer to a new basis or grid without changing what it computes: tabulate the old edge functions on samples and solve a least squares problem per input in the new basis. Grid update, grid extension (refinement) and conversion between bases (e.g. B-spline -> Chebyshev) are all this one operation.

adaptive_grid(x, n_points, margin=0.0, grid_eps=0.02)

Per-input grids (in, n_points) from samples x of shape (N, in).

resample_grid(grid, n_points)

Interpolate grids (in, n) to n_points points, keeping their shape.

grid_ranges(basis, params)

Per-input (lo, hi) of the basis grid, (in, 2), or None without a grid.

sample_inputs(ranges, n)

n uniform points per input inside ranges (in, 2): (n, in).

refit(old_basis, old_params, old_coef, new_basis, new_params, x, ridge=0.0)

Coefficients (out, in, new_basis.size) reproducing the old edge functions.

x (N, in) are the points where they must match; params are expanded per input. Solved per input with lstsq (minimum-norm when a basis function has no samples), so a new basis that contains the old one reproduces it to rounding error. ridge > 0 adds Tikhonov damping, relative to the mean squared basis value.

kanforge.regularization

Sparsity regularisation of KAN layers, written once for every backend.

pykan's penalty works on edge magnitudes |phi_oi| averaged over a batch: L1 of the magnitudes plus the entropy of their row and column distributions, which pushes a layer towards few important edges (needed before pruning and symbolic regression). Optional terms penalise the coefficients themselves and their differences (smoothness). Without data, efficient-kan's variant uses mean |coef| as edge magnitude.

penalty(magnitude, coef, *, l1=1.0, entropy=2.0, coef_l1=0.0, coef_smooth=0.0, eps=1e-08)

Penalty of one layer from edge magnitudes (out, in) and coefficients (out, in, n).

edge_magnitude(activations)

mean |phi_oi(x)| over all batch dimensions: (..., out, in) -> (out, in).

kanforge.flops

FLOPs and parameters of KAN and MLP models, for comparisons at equal cost.

Counts are per sample, forward pass, following KANbeFair's convention: a multiply-add is 2 FLOPs, elementwise arithmetic 1, a transcendental function (exp, tanh, sin, ...) NONLINEAR FLOPs (the same constant for KANs and MLPs, so neither side is favoured). Basis evaluation is counted from the algorithm the plain-ops path runs.

basis_flops(basis)

FLOPs to evaluate all functions of basis at one input value.

count(model)

{"flops": per-sample forward FLOPs, "params": trainable parameters}.

Handles KAN layers and models of every backend, and torch nn.Linear / activations / nn.Sequential MLPs for the comparison side.

kanforge.kernels.plan

What a fused kernel has to compute for a basis, independent of the backend.

A KernelPlan describes the per-input kernel features phi_j(x) that the generated code evaluates in registers, and how derivatives are obtained:

• cheb and fourier are folded: the m-th x-derivative of every basis function is a fixed linear combination of the kernel features (phi_k^(m) = sum_j R_m[k, j] T_j), so kernels never differentiate and R_m is applied to the coefficients outside the kernel. All polynomial families with a tanh or identity input map reduce to Chebyshev features.
• bspline and bump are native: the derivative order is compiled into the kernel.

KernelPlan(kind, consts, size, fold_key=(), grid=False) dataclass

n_features(m)

Kernel features evaluated for derivative order m.

fold(m)

R_m (size, n_features(m)) for folded plans, None otherwise.

bump_poly(m)

Polynomial (in u, or in tanh(u)) of the m-th derivative of a bump.

plan_for(basis)

The kernel plan of basis or None when it needs the plain-ops path (trainable basis parameters or a basis without a kernel yet).