build_grid

kan_module.build_grid

Deterministic knot grid [in_features, grid_size + 2*spline_order + 1] (NF kan.py:42-51).

Parameter	Type	Default	Description
in_features	int	-	(undocumented)
grid_size	int	-	(undocumented)
spline_order	int	-	(undocumented)
grid_range	tuple[float, float]	-	(undocumented)

Returns: jnp.ndarray

b_splines

kan_module.b_splines

Cox-de Boor B-spline bases. x:[B, in], grid:[in, G] -> [B, in, grid_size+spline_order].

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
grid	jnp.ndarray	-	(undocumented)
spline_order	int	-	(undocumented)

Returns: jnp.ndarray

KANLinear

kan_module.KANLinear · inherits nnx.Module

A KAN edge layer: base = Linear(SiLU(x), base_weight); spline = Linear(b_splines(x), spline*scaler).

__init__(self, in_features, out_features, *, grid_size, spline_order, scale_noise, scale_base, scale_spline, enable_standalone_scale_spline, grid_range, rngs: nnx.Rngs)

Parameter	Type	Default	Description
in_features	-	-	(undocumented)
out_features	-	-	(undocumented)
grid_size	-	-	(undocumented)
spline_order	-	-	(undocumented)
scale_noise	-	-	(undocumented)
scale_base	-	-	(undocumented)
scale_spline	-	-	(undocumented)
enable_standalone_scale_spline	-	-	(undocumented)
grid_range	-	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray) -> jnp.ndarray

x: [B, in_features] -> [B, out_features].

Parameter	Type	Default	Description
x	jnp.ndarray	-	[B, in_features]

Returns: jnp.ndarray

KANNet

kan_module.KANNet · inherits nnx.Module

Stack of KANLinear layers: [input_size -> hidden_size (xN) -> h]. I/O [B, L, 1] -> [B, h, 1].

__init__(self, *, h, input_size, n_hidden_layers, hidden_size, grid_size, spline_order, scale_noise, scale_base, scale_spline, enable_standalone_scale_spline, grid_range, rngs: nnx.Rngs)

Parameter	Type	Default	Description
h	-	-	(undocumented)
input_size	-	-	(undocumented)
n_hidden_layers	-	-	(undocumented)
hidden_size	-	-	(undocumented)
grid_size	-	-	(undocumented)
spline_order	-	-	(undocumented)
scale_noise	-	-	(undocumented)
scale_base	-	-	(undocumented)
scale_spline	-	-	(undocumented)
enable_standalone_scale_spline	-	-	(undocumented)
grid_range	-	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray) -> jnp.ndarray

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)

Returns: jnp.ndarray ([B, h, 1].)