Skip to content

Tensor Network Kernel Machines (TNKM)

TNKM is a JAX-based library for machine learning with tensor-network parameterizations.

It provides kernel machines built from low-rank tensor decompositions, enabling scalable models with explicit rank control and structured inductive bias.

TNKM is designed for researchers working on structured function approximation, system identification, and time-series modeling.

Why TNKM?

  • Structured kernel models with controllable complexity via tensor ranks
  • CP and TT tensor-network architectures
  • Two training paradigms
    • Alternating Least Squares (ALS)
    • Gradient-based optimization (JAX + Optax)
  • JAX-native implementation for composability, speed and hardware acceleration

Quick example

Install

TNKM is available on PyPI. See the installation guide for details.

pip install tnkm

Fit a CP-based tensor-network kernel machine with polynomial features:

import jax.numpy as jnp

from tnkm.models import CPKM
from tnkm.optim import train_als
from tnkm.features import ProductFeatures, PolyFeature

d_dim = 2
x = jnp.ones((32, d_dim))
y = jnp.ones((32,))

features = ProductFeatures((PolyFeature(4, k_col=i) for i in range(d_dim)))
model = CPKM(features, rank=2, seed=0)

train_als(model, x, y, n_epoch=3, gamma_w=1e-3, beta_e=1.0)
y_pred = model.predict(x)
\[ \begin{split} &\min_{\boldsymbol{v}}\dfrac{\beta}{2} \|\boldsymbol{y} -\boldsymbol{\Phi}\boldsymbol{g}(\boldsymbol{v})\|_2^2 + \dfrac{\gamma}{2}\|\boldsymbol{v}\|_2^2 \\ & \text{s.t. }\boldsymbol{g}(\boldsymbol{v}) = \sum_{r=1}^{R} \boldsymbol{v}^{(D)}_{r} \otimes \dots \otimes \boldsymbol{v}^{(1)}_{r}. \end{split} \]

Next steps

Getting started

Start with the installation guide, then continue with the quickstart tutorial to build your first tensor-network kernel machine.

Examples

Ready-to-run workflows and use cases:

Citation

TNKM is an open-source research software project. If you use TNKM in academic work, please cite the software:

Albert Saiapin and Kim Batselier. (2026). tnkm: A JAX Framework for Tensor Network Kernel Machines. Zenodo. https://doi.org/10.5281/zenodo.21841905

For the scientific background and methodology, please also cite the associated research paper:

Albert Saiapin and Kim Batselier. (2026). Tensor Network Kernel Machines: A JAX Framework for Machine Learning and Nonlinear System Identification. https://arxiv.org/abs/2608.07043

If you find the project useful, consider starring the repository on GitHub.