Tensor Network Kernel Machines (TNKM)

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.

Features

• 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

Installation

From PyPI

pip install tnkm

From source

git clone https://github.com/AlbMLpy/tnkm.git
cd tnkm
pip install .

Quick example

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)

Documentation

Available at https://albmlpy.github.io/tnkm.

Contributing

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.

Citation

TNKM is an open-source research software project. If you use TNKM in academic work, please cite the specific software version used. Citation information will be updated when a DOI or publication becomes available.

(Also consider starring the project on GitHub)

License

MIT License