KAN

chronax.models.KAN · inherits BaseForecaster

Univariate Kolmogorov-Arnold Network forecaster (JAX/Flax-NNX port of neuralforecast.KAN). An MLP whose layers learn per-edge B-spline activations.

Maintenance status: Active univariate forecaster. Integrates with the BaseForecaster interface, including conformal prediction intervals via predict(level=...), pickle round-trip, and forecast(fitted=True).

scaler selects the per-window input scaling: "identity" (default, matches neuralforecast — splines inactive on raw-scale data) or "robust" (median/MAD; normalizes into the spline grid so the B-splines engage). float32 throughout.

Attributes:

Name	Type	Description
uses_exog	bool	False
alias	str	KAN
conformal_params	None	(undocumented)
model_	KANNet or None	The fitted network model.

__init__(self, h: int, input_size: int = -1, grid_size: int = 5, spline_order: int = 3, scale_noise: float = 0.1, scale_base: float = 1.0, scale_spline: float = 1.0, enable_standalone_scale_spline: bool = True, grid_eps: float = 0.02, grid_range: tuple = (-1.0, 1.0), n_hidden_layers: int = 1, hidden_size: int = 512, max_steps: int = 1000, learning_rate: Union[float, Callable[[int], float]] = 1e-3, windows_batch_size: int = 1024, loss: Union[str, LossFn] = "mae", scaler: Union[str, Scaler] = "identity", random_seed: int = 1, alias: str = "KAN")

(undocumented)

Parameter	Type	Default	Description
h	int	-	(undocumented)
input_size	int	-1	(undocumented)
grid_size	int	5	(undocumented)
spline_order	int	3	(undocumented)
scale_noise	float	0.1	(undocumented)
scale_base	float	1.0	(undocumented)
scale_spline	float	1.0	(undocumented)
enable_standalone_scale_spline	bool	True	(undocumented)
grid_eps	float	0.02	(undocumented)
grid_range	tuple	(-1.0, 1.0)	(undocumented)
n_hidden_layers	int	1	(undocumented)
hidden_size	int	512	(undocumented)
max_steps	int	1000	(undocumented)
learning_rate	Union[float, Callable[[int], float]]	1e-3	(undocumented)
windows_batch_size	int	1024	(undocumented)
loss	Union[str, LossFn]	"mae"	(undocumented)
scaler	Union[str, Scaler]	"identity"	(undocumented)
random_seed	int	1	(undocumented)
alias	str	"KAN"	(undocumented)

fit(self, y: jnp.ndarray, X: jnp.ndarray | None = None) -> Self

Fit on a 1-D series. Raises NotImplementedError on exog; ValueError if y is not 1-D or shorter than input_size+h; RuntimeError on divergence.

Parameters:

Parameter	Type	Default	Description
y	jnp.ndarray	-	(undocumented)
X	jnp.ndarray or None	None	(undocumented)

Returns: Self (the fitted forecaster; sets self.model_). Raises: * NotImplementedError * ValueError * RuntimeError

predict(self, h: int, X: jnp.ndarray | None = None, level: list[int | float] | None = None) -> dict

Forecast h steps (1 <= h <= self.h). level -> inherited conformal path.

Parameters:

Parameter	Type	Default	Description
h	int	-	(undocumented)
X	jnp.ndarray or None	None	(undocumented)
level	list[int | float] or None	None	(undocumented)

Returns: dict (Keys include {"mean": jnp.ndarray}. Additional keys for prediction intervals are added if level is provided.) Raises: * RuntimeError * ValueError

forecast(self, y: jnp.ndarray, h: int, X: jnp.ndarray | None = None, X_future: jnp.ndarray | None = None, level: list[int | float] | None = None, fitted: bool = False) -> dict

Stateless fit-then-predict; fitted=True adds one-step-ahead fitted values.

Parameters:

Parameter	Type	Default	Description
y	jnp.ndarray	-	(undocumented)
h	int	-	(undocumented)
X	jnp.ndarray or None	None	(undocumented)
X_future	jnp.ndarray or None	None	(undocumented)
level	list[int | float] or None	None	(undocumented)
fitted	bool	False	(undocumented)

Returns: dict (Keys include {"mean": jnp.ndarray}. If fitted=True, includes "fitted": jnp.ndarray.) Raises: * NotImplementedError