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.0.4):

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/

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

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