Optimizers

optim

Optimizers for training TNKM models using JAX.

train_als

train_als(
    tnkm: TNKM,
    x: Array,
    y: Array,
    n_epoch: int,
    gamma_w: float = 0.0,
    beta_e: float = 1.0,
    d_core_stop: int | None = None,
    reg_mode: str = "cores",
    tracker: TrackerLike | None = None,
) -> None

Train a TNKM model using Alternating Least Squares (ALS).

ALS optimizes tensor-network parameters by sequentially updating one core at a time while keeping all others fixed. Each update is obtained by solving a regularized least-squares problem.

Parameters:

• tnkm (TNKM) –

  Tensor network kernel machine model, either CP- or TT-based.

• x (Array) –

  Input data of shape (n_samples, d_dim).

• y (Array) –

  Target values of shape (n_samples,).

• n_epoch (int) –

  Number of training epochs. Must be positive.

• gamma_w (float, default: 0.0 ) –

  Weight of the L2 regularization term.

• beta_e (float, default: 1.0 ) –

  Weight of the mean squared error term.

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

  If provided, the final epoch updates only cores with indices smaller or equal to d_core_stop. This can be used to freeze the remaining cores during the final sweep.

• reg_mode (str, default: 'cores' ) –

  Regularization strategy.

  • "cores": apply L2 regularization directly to the tensor cores.
  • "tensor": apply L2 regularization to the reconstructed tensor.
• tracker (TrackerLike | None, default: None ) –

  Optional training callback invoked before training and after each epoch.

Notes

Convergence depends on initialization and regularization strength.

Raises:

• ValueError –

  If n_epoch, gamma_w, beta_e, d_core_stop, or reg_mode has an invalid value.

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. Section 3.1.

Examples:

>>> import jax.numpy as jnp
>>> from tnkm.models import CPKM
>>> from tnkm.optim import train_als
>>> 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)
>>> x = jnp.arange(6).reshape((3, 2))
>>> y = x.sum(axis=1)
>>> tnkm.predict(x) # untrained model output
Array([-2.0700777, -8.881002 , 38.90358  ], dtype=float32)
>>> train_als(tnkm, x, y, n_epoch=1, gamma_w=0.001, beta_e=1.0)
>>> tnkm.predict(x) # array is close to y
Array([1.003613, 4.995645, 9.0013  ], dtype=float32)

train_grad

train_grad(
    model: ModelLike,
    train_loader: Iterable[Batch],
    solver: SolverLike,
    n_epoch: int,
    loss_f: Callable[[Any, Batch], Array],
    tracker: TrackerLike | None = None,
    holomorphic: bool = False,
) -> list[float]

Train a model with a first-order optimizer.

This function performs gradient-based optimization of model parameters using an Optax optimizer over multiple epochs.

Parameters:

• model (ModelLike) –

  Model instance containing trainable parameters. Must expose a mutable params attribute (e.g., a JAX PyTree).

• train_loader (Iterable[Batch]) –

  Iterable of training batches (X, y).

• solver (SolverLike) –

  Optax optimizer (e.g., optax.sgd, optax.adam).

• n_epoch (int) –

  Number of training epochs. Must be positive.

• loss_f (Callable[[Any, Batch], Array]) –

  Loss function mapping (params, batch) → scalar loss.

• tracker (TrackerLike | None, default: None ) –

  Optional callback for monitoring training progress.

• holomorphic (bool, default: False ) –

  If True, assumes loss_f is holomorphic and uses complex-valued differentiation.

Returns:

• list[float] –

  Mean training loss for each epoch.

Raises:

• ValueError –

  If n_epoch is not positive.

• TypeError –

  If loss_f is not callable.

Examples:

>>> import optax
>>> import jax.numpy as jnp
>>> from tnkm.models import CPKM
>>> from tnkm.optim import train_grad
>>> from tnkm.dataloader import DataLoader
>>> 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)
>>> loss_f = tnkm.make_loss(gamma_w=0.001, beta_e=1.0, reg_mode="cores")
>>> x = jnp.arange(6).reshape((3, 2))
>>> y = x.sum(axis=1)
>>> tnkm.predict(x) # untrained model output
Array([-2.0700777, -8.881002 , 38.90358  ], dtype=float32)
>>> train_loader = DataLoader(x, y, batch_size=1, seed=0)
>>> solver = optax.adam(learning_rate=1e-1)
>>> losses = train_grad(tnkm, train_loader, solver, 200, loss_f)
>>> tnkm.predict(x) # array is close to y
Array([1.064034, 4.974135, 9.003395], dtype=float32)