Tensor Networks for Machine Learning

tn4ml is a Python library that handles tensor networks for machine learning applications.
It is built on top of Quimb, for Tensor Network objects, and JAX, for optimization pipeline.
For now, the library supports 1D Tensor Network structures:

• Matrix Product State
• Matrix Product Operator
• Spaced Matrix Product Operator

It supports different embedding functions, initialization techniques, objective functions and optimization strategies.

Installation

First create a virtualenv using pyenv or conda. Then install the package and its dependencies.

With pip (tag v1.1.0):

pip install tn4ml

or directly from github:

pip install -U git+https://github.com/bsc-quantic/tn4ml.git

If you want to test and edit the code, you can clone the local version of the package and install it.

git clone https://github.com/bsc-quantic/tn4ml.git
pip install -e tn4ml/

If you want to install dependices for docs, test and examples:

pip install "tn4ml[docs]"

pip install "tn4ml[test]"

pip install "tn4ml[examples]"

Developer checks

Install the pre-commit hooks before contributing:

pip install pre-commit
pre-commit install

Run all hooks manually:

pre-commit run --all-files

The hooks run fast local checks from CI/CD (ruff, mypy, bandit) and strip notebook outputs/local metadata with nbstripout, including notebook kernel display names such as local environment names.

Accelerated runtime

(Optional) To improve runtime precision set these flags:

jax.config.update("jax_enable_x64", True)
jax.config.update('jax_default_matmul_precision', 'highest')

Running on GPU
Before everything install JAX version that supports CUDA and its suitable for runs on GPU.
Checkout how to install here: jax[cuda]

Next, at the beginning of your script set:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # Use GPU 0 - or set any GPU ID
import jax
jax.config.update("jax_platform_name", 'gpu')

Then when training Model set:

device = 'gpu'
model.configure(device=device)

Documentation

Visit tn4ml.readthedocs.io

Example notebooks

TN for Classification
TN for Anomaly Detection
TN for Anomaly Detection with DMRG-like method

Examples from the paper

Breast Cancer Classification
Unsupervised Learning with MNIST
MPS for Anomaly Detection in the Latent Space of Proton Collision Events at the LHC

Citation

If you use tn4ml in your work, please cite the following paper: arXiv:2502.13090

@article{puljak2025tn4mltensornetworktraining,
      title={tn4ml: Tensor Network Training and Customization for Machine Learning},
      author={Ema Puljak and Sergio Sanchez-Ramirez and Sergi Masot-Llima and Jofre Vallès-Muns and Artur Garcia-Saez and Maurizio Pierini},
      year={2025},
      eprint={2502.13090},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2502.13090},
      }

License

MIT license - check it out here