TNKM is a JAX-based library for machine learning with tensor-network parameterizations.
It combines kernel methods with low-rank tensor decompositions to build scalable models with explicit rank control.
TNKM is intended for research in machine learning, system identification, and time-series modeling.
- CP and TT tensor-network kernel machines
- Polynomial, Fourier, B-spline and other feature maps
- Alternating Least Squares (ALS) and gradient-based optimization (Optax)
- Explicit rank control for balancing model complexity and expressiveness
- JAX-native implementation with JIT compilation and hardware acceleration
pip install tnkmgit clone https://github.com/AlbMLpy/tnkm.git
cd tnkm
pip install .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)Available at https://albmlpy.github.io/tnkm.
Contributions to tnkm are welcome. If you would like to report issues,
suggest improvements, or contribute code, please see
CONTRIBUTING.md for development guidelines,
installation instructions, and information about the contribution workflow.
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, which will be linked here once available.
If you find the project useful, consider starring the repository on GitHub.
MIT License