MLX¶
kanforge.mlx.layers
¶
MLX KAN layers for Apple silicon.
Same parameters and math as kanforge.torch; custom Metal kernels will
plug in behind these modules later. Buffers (grids) are frozen arrays.
KANLinear(in_features, out_features, basis='bspline', *, bias=True, base_activation=nn.silu, groups=None, edge_scale=False, grid_ranges=None, coef_init='normal', coef_init_std=0.1, edge_dropout=0.0, backend='auto')
¶
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.
Parameter names and shapes match kanforge.torch.KANLinear.
backend as in torch: "kernel" runs the fused Metal kernel on the GPU.
uses_kernel()
¶
Whether forward runs the fused Metal kernel (default device is the GPU).
edge_derivative(x, order=1)
¶
d^order phi_oi / dx_i^order for every edge, shape (..., out, in).
from_linear(linear, activation=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)
¶
KAN(widths, basis='bspline', grid_ranges=None, **layer_kwargs)
¶
Bases: Module
Stack of KANLinear layers, e.g. KAN([2, 5, 1], basis="chebyshev").
MultKANLinear(in_features, n_sum, n_mult=0, basis='bspline', *, arity=2, **kwargs)
¶
Bases: Module
KAN layer with addition and multiplication nodes (see kanforge.torch).
MultKAN(widths, basis='bspline', *, arity=2, **layer_kwargs)
¶
Bases: Module
MultKAN([2, (2, 1), 1]) (see kanforge.torch.MultKAN).
from_mlp(mlp, basis='bspline', **kwargs)
¶
MLX version of kanforge.torch.from_mlp.
rebuild_kan(layers)
¶
A KAN made of existing layers (used by pruning).
layer_inputs(model, x)
¶
(KANLinear, its input) for KANLinear, MultKANLinear, KAN, MultKAN, nn.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.mlx.conv
¶
KAN convolutions for MLX (channels-last, like mlx.nn.ConvNd).
Same math and patch-feature order (C_in/groups, k_1, ..., k_n) as
kanforge.torch.conv; inputs are (N, *spatial, C).
kanforge.mlx.blocks
¶
MLX versions of the architecture blocks: transformer FFN, graph layers,
recurrent cells and scientific ML tools. Same math and arguments as the
corresponding kanforge.torch modules (transformer, graph,
recurrent, science); see their docstrings.
derivative(fn, x, order=1, dim=0)
¶
d^order u / dx_dim^order for u = fn(x) (functional, via mx.grad).