Kolmogorov-Arnold Network Helpers

KAN layers replace each scalar edge by a small trainable one-dimensional function. TorchLean keeps that structure visible: an edge family first expands every scalar input into basis features, and the KAN layer learns one coefficient per (output, input, basis) edge.

The first built-in family uses triangular piecewise-linear hats. Users can add another family by constructing KANEdgeFamily: provide a basis dimension and a TorchLean model that maps Vec inDim to Vec (inDim * basisDim).

References:

• Z. Liu et al., "KAN: Kolmogorov-Arnold Networks", arXiv:2404.19756.
• C. de Boor, "A Practical Guide to Splines", Springer, 1978/2001.

@[reducible, inline]

abbrev NN.API.nn.models.KANShape.vec (n : ℕ) : Spec.Shape

Unbatched vector shape used by KAN edge bases.

@[reducible, inline]

abbrev NN.API.nn.models.KANShape.mat (rows cols : ℕ) : Spec.Shape

Matrix shape used by batched KAN models and basis tables.

structure NN.API.nn.models.KANEdgeFamily : Type 2

Backend-compatible KAN edge family.

An edge family turns each scalar input coordinate into basisDim features. A KAN layer then applies a learned linear map to all expanded features. The basis is a TorchLean model fragment, not an arbitrary Lean callback, so the resulting KAN can run in eager, compiled, CPU, and CUDA training paths supported by the underlying operations.

• name : String

  Short label shown in model summaries and training metadata.

• basisDim : ℕ

  Number of basis features produced per scalar input coordinate.

• basis (inDim : ℕ) : Sequential (KANShape.vec inDim) (KANShape.vec (inDim * self.basisDim))

  Basis expansion for an unbatched vector of length inDim.

@[reducible, inline]

abbrev NN.API.nn.models.KANEdgeFamily.basisShape (edge : KANEdgeFamily) (inDim : ℕ) : Shape

Shape of the edge-basis expansion for inDim scalar inputs.

structure NN.API.nn.models.KANPiecewiseLinear : Type

Configuration for triangular piecewise-linear KAN edge bases.

The basis functions are hats centered at the integer knots 0, ..., gridSize - 1. The input is multiplied by inputScale before the hats are evaluated. For normalized data in [0, 1], setting inputScale = gridSize - 1 spreads the grid across the full interval.

• gridSize : ℕ

  Number of knots, hence the number of basis functions per scalar coordinate.

• inputScale : ℕ

  Scale applied before basis evaluation; use gridSize - 1 for normalized [0, 1] inputs.

@[implicit_reducible]

instance NN.API.nn.models.instReprKANPiecewiseLinear : Repr KANPiecewiseLinear

def NN.API.nn.models.instReprKANPiecewiseLinear.repr : KANPiecewiseLinear → ℕ → Std.Format

def NN.API.nn.models.KANPiecewiseLinear.basisLayer (cfg : KANPiecewiseLinear) (inDim : ℕ) : Sequential (KANShape.vec inDim) (KANShape.vec (inDim * cfg.gridSize))

Expand x : Vec inDim to all triangular basis features.

The output is flattened row-major from a (gridSize × inDim) table: [basis_0(x_0), ..., basis_0(x_n), basis_1(x_0), ...].

Each basis value is relu(1 - |inputScale * x_i - k|), expressed directly in the ordinary TorchLean op language rather than through an opaque spline evaluator.

def NN.API.nn.models.KANPiecewiseLinear.edgeFamily (cfg : KANPiecewiseLinear) : KANEdgeFamily

Turn piecewise-linear triangular bases into a general KAN edge family.

structure NN.API.nn.models.KANConfig : Type 2

Configuration for a KAN over batched row vectors.

• batch : ℕ

  Leading minibatch dimension.

• inDim : ℕ

  Number of scalar input coordinates.

• hidden : List ℕ

  Hidden KAN widths. Each entry creates one KAN layer followed by tanh.

• outDim : ℕ

  Number of output coordinates/classes.

• edge : KANEdgeFamily

  Edge basis family. The default is a compact triangular piecewise-linear basis.

• seedBase : ℕ

  Base seed used for learned edge coefficients and biases.

@[reducible, inline]

abbrev NN.API.nn.models.kanInShape (cfg : KANConfig) : Shape

Input shape (batch × inDim) for a KAN config.

@[reducible, inline]

abbrev NN.API.nn.models.kanOutShape (cfg : KANConfig) : Shape

Output shape (batch × outDim) for a KAN config.

def NN.API.nn.models.kanLayer (inDim outDim : ℕ) (edge : KANEdgeFamily) (seedW seedB : ℕ := 0) : Sequential (KANShape.vec inDim) (KANShape.vec outDim)

One unbatched KAN layer.

The layer first applies the selected edge basis to every input coordinate, then learns coefficients with an ordinary linear map from the expanded features to outDim.

def NN.API.nn.models.kanGo (edge : KANEdgeFamily) (inDim : ℕ) (hidden : List ℕ) (outDim seed : ℕ) : Sequential (KANShape.vec inDim) (KANShape.vec outDim)

Recursive unbatched KAN stack. Hidden layers use tanh; the final layer is linear in bases.

def NN.API.nn.models.KAN (cfg : KANConfig) : M (Sequential (kanInShape cfg) (kanOutShape cfg))

Build a batched KAN model.

Task semantics are deliberately not baked into the model name: use Trainer.new with task := .regression, .classification, .crossEntropy, or .custom ... with the same KAN constructor.