Models

models

Tensor Network Kernel Machine models.

TNKM

Bases: ABC

Abstract base class for tensor-network kernel machine models.

Concrete implementations must provide prediction and loss-construction methods and expose mutable model parameters through the params property.

params abstractmethod property writable

params: Params

Return model parameters.

name abstractmethod property

name: str

Return model identifier string.

predict abstractmethod

predict(x: ArrayLike) -> Array

Compute model predictions.

make_loss abstractmethod

make_loss(
    *args, **kwargs
) -> Callable[[Params, Batch], Array]

Create a model-specific loss function.

make_predict abstractmethod

make_predict() -> Callable[[Params, Array], Array]

Create a model-specific prediction function.

CPKM dataclass

CPKM(
    fmap: ProductFeatures,
    rank: int,
    seed: int | None = None,
    dtype: dtype = jnp.float32,
)

Bases: TNKM

CP tensor-network kernel machine.

Implements a tensor-network kernel machine based on the Canonical Polyadic (CP) decomposition. The model combines a product feature map with a trainable CP tensor network to efficiently represent high-dimensional functions.

Parameters:

• fmap (ProductFeatures) –

  Product feature map used to construct the input feature representation.

• rank (int) –

  CP rank (number of rank-1 tensor components). Must be positive.

• seed (int | None, default: None ) –

  Random seed for reproducible parameter initialization.

• dtype (dtype, default: jnp.float32 ) –

  Data type of the model parameters and computations.

Raises:

• ValueError –

  If rank is not a positive integer.

• TypeError –

  If seed is not None or an integer, or if fmap is not an instance of ProductFeatures.

References

[1] F. Wesel, K. Batselier, "Large-Scale Learning with Fourier Features and Tensor Decompositions", Advances in Neural Information Processing Systems, 2021.

[2] A. Saiapin, K. Batselier, "Laplace Approximation for Bayesian Tensor Network Kernel Machines", 2026.

Examples:

>>> from tnkm.models import CPKM
>>> from tnkm.features import PolyFeature, ProductFeatures
>>> fmap = ProductFeatures((PolyFeature(i+2, k_col=i)) for i in range(2))
>>> tnkm = CPKM(fmap, rank=4, seed=1)
>>> tnkm.name
'cpkm'
>>> len(tnkm.params)
2
>>> tnkm.params[0].shape
(2, 4)
>>> tnkm.params[1].shape
(3, 4)

params property writable

params: Params

Return the CPKM model parameters.

These are the CP tensor cores stored in the underlying tensor network.

name property

name: str

CPKM model identifier.

predict

predict(x: ArrayLike) -> Array

Predict model outputs for the given inputs.

Parameters:

• x (ArrayLike) –

  Input data of shape (n_samples, ...).

Returns:

• Array –

  Model predictions of shape (n_samples,).

make_loss

make_loss(
    gamma_w: float, beta_e: float, reg_mode: str
) -> Callable[[Params, Batch], Array]

Construct a JAX-jittable L2-regularized MSE loss function for CPKM.

Parameters:

• gamma_w (float) –

  Weight of the regularization term.

• beta_e (float) –

  Weight of the data term (MSE).

• reg_mode (str) –

  Type of regularization to use:

  • 'cores' : squared L2 norm over CP cores
  • 'tensor' : full tensor L2 regularization

Returns:

• Callable[[Params, Batch], Array] –

  Loss function loss(params, batch) returning a scalar value.

make_predict

make_predict() -> Callable[[Params, Array], Array]

Create a JAX-jittable prediction function.

Returns:

• Callable[[Params, Array], Array] –

  Function predict(params, x) that maps inputs x to model outputs using the given parameters. The output has shape (n_samples, ...).

TTKM dataclass

TTKM(
    fmap: ProductFeatures,
    rank: int | Sequence[int],
    seed: int | None = None,
    dtype: dtype = jnp.float32,
)

Bases: TNKM

TT tensor-network kernel machine.

Implements a tensor-network kernel machine based on the Tensor Train (TT) decomposition. The model combines a product feature map with a trainable TT tensor network to efficiently represent high-dimensional functions.

Parameters:

• fmap (ProductFeatures) –

  Product feature map used to construct the input feature representation.

• rank (int | Sequence[int]) –

  Specification of TT ranks.

  • int: uniform internal rank ([1, r, ..., r, 1])
  • sequence: full TT rank vector of length d_dim + 1 with boundary constraints r_0 = r_d = 1
• seed (int | None, default: None ) –

  Random seed for reproducible parameter initialization.

• dtype (dtype, default: jnp.float32 ) –

  Data type of the model parameters and computations.

Raises:

• ValueError –

  If rank is not a positive integer, if a rank sequence is empty, if rank[0] != 1 or rank[-1] != 1, if any TT rank is not a positive integer, if rank is neither an int nor a sequence of ints, or if len(fmap.local_dims) <= 1.

• TypeError –

  If seed is not None or an integer, or if fmap is not an instance of ProductFeatures.

References

[1] R. Karagoz, K. Batselier, "Nonlinear system identification with regularized Tensor Network B-splines", Automatica, 2020.

[2] A. Saiapin, K. Batselier, "Laplace Approximation For Tensor Train Kernel Machines In System Identification", 2025.

Examples:

>>> from tnkm.models import TTKM
>>> from tnkm.features import PolyFeature, ProductFeatures
>>> fmap = ProductFeatures((PolyFeature(i+2, k_col=i)) for i in range(2))
>>> tnkm = TTKM(fmap, rank=4, seed=1)
>>> tnkm.name
'ttkm'
>>> len(tnkm.params)
2
>>> tnkm.params[0].shape
(1, 2, 4)
>>> tnkm.params[1].shape
(4, 3, 1)

params property writable

params: Params

Return the TTKM model parameters.

These are the TT tensor cores stored in the underlying tensor network.

name property

name: str

TTKM model identifier.

predict

predict(x: ArrayLike) -> Array

Predict model outputs for the given inputs.

Parameters:

• x (ArrayLike) –

  Input data of shape (n_samples, ...).

Returns:

• Array –

  Model predictions of shape (n_samples,).

make_loss

make_loss(
    gamma_w: float, beta_e: float, reg_mode: str
) -> Callable[[Params, Batch], Array]

Construct a JAX-jittable L2-regularized MSE loss function for TTKM.

Parameters:

• gamma_w (float) –

  Weight of the regularization term.

• beta_e (float) –

  Weight of the data term (MSE).

• reg_mode (str) –

  Type of regularization to use:

  • 'cores' : squared L2 norm over TT cores
  • 'tensor' : full tensor L2 regularization

Returns:

• Callable[[Params, Batch], Array] –

  Loss function loss(params, batch) returning a scalar value.

make_predict

make_predict() -> Callable[[Params, Array], Array]

Create a JAX-jittable prediction function.

Returns:

• Callable[[Params, Array], Array] –

  Function predict(params, x) that maps inputs x to model outputs using the given parameters. The output has shape (n_samples, ...).