torchTT

Tensor-Train decomposition in pytorch

Tensor-Train decomposition package written in Python on top of pytorch. Supports GPU acceleration and automatic differentiation. It also contains routines for solving linear systems in the TT format and performing adaptive cross approximation (the AMEN solver/cross interpolation is inspired form the MATLAB TT-Toolbox). Some routines are implemented in C++ for an increased execution speed.

Installation

Requirements

Following requirements are needed:

• python>=3.8
• torch>=1.7.0
• numpy>=1.18
• opt_einsum

The GPU (if available) version of pytorch is recommended to be installed. Read the official installation guide for further info.

Using pip

You can install the package using the pip command:

pip install torchTT

The latest github version can be installed using:

pip install git+https://github.com/ion-g-ion/torchTT

One can also clone the repository and manually install the package:

git clone https://github.com/ion-g-ion/torchTT
cd torchTT
python setup.py install

Using uv

You can install the package using uv:

uv pip install torchTT

The latest github version can be installed using:

uv pip install git+https://github.com/ion-g-ion/torchTT

One can also clone the repository and install the package using uv:

git clone https://github.com/ion-g-ion/torchTT
cd torchTT
uv sync

Or install in editable mode:

uv pip install -e .

Development Installation

For development, you may want to install the package with additional development dependencies (pytest, sphinx, ipykernel, matplotlib):

Using pip:

pip install -e ".[dev]"

Using uv:

uv sync --extra dev

or

uv pip install -e ".[dev]"

This will install the package in editable mode along with all development tools needed for testing, building documentation, and working with Jupyter notebooks.

Components

The main modules/submodules that can be accessed after importing torchtt are briefly desctibed in the following table. Detailed description can be found here.

Component	Description
torchtt	Basic TT class and basic linear algebra functions.
torchtt.solvers	Implementation of the AMEN solver.
torchtt.grad	Wrapper for automatic differentiation.
torchtt.manifold	Riemannian gradient and projection onto manifolds of tensors with fixed TT rank.
torchtt.nn	Basic TT neural network layer.
torchtt.interpolate	Cross approximation routines.

Tests

The directory tests/ from the root folder contains all the unittests. To run them use the command:

pytest tests/

Documentation and examples

The documentation can be found here. Following example scripts (as well as python notebooks) are also provied provided as part of the documentation:

• basic_tutorial.py / basic_tutorial.ipynp: This contains a basic tutorial on decomposing full tensors in the TT format as well as performing rank rounding, slicing (Try on Google Colab).
• basic_linalg.py / basic_linalg.ipynp: This tutorial presents all the algebra operations that can be performed in the TT format (Try on Google Colab).
• efficient_linalg.py / efficient_linalg.ipynb: contains the DMRG for fast matves and AMEN for elementwise inversion in the TT format (Try on Google Colab).
• automatic_differentiation.py / automatic_differentiation.ipynp: Basic tutorial on AD in torchtt (Try on Google Colab).
• cross_interpolation.py / cross_interpolation.ipynb: In this script, the cross interpolation emthod is exemplified (Try on Google Colab).
• system_solvers.py / system_solvers.ipynb: This contains the bais ussage of the multilinear solvers (Try on Google Colab).
• gpu_acceleration.py / gpu_acceleration.ipynb: This provides an example on how to use the GPU acceleration (Try on Google Colab).
• basic_nn.py / basic_nn.ipynb: This provides an example on how to use the TT neural network layers (Try on Google Colab).
• mnist_nn.py / mnist_nn.ipynb: Example of TT layers used for image classification (Try on Google Colab).
• manifold.py / manifold.ipynb: This demonstrates Riemannian gradient descent on manifolds of tensors with fixed TT rank (Try on Google Colab).
• random_tt.py: This script shows how to generate random TT tensors with different variances (Try on Google Colab).
• tensor_completion.py: This example demonstrates tensor completion using manifold learning with Riemannian gradient descent (Try on Google Colab).

Building Documentation

The documentation is generated using sphinx. To build it locally, you need:

1. Install development dependencies (see Development Installation above)

2. Install pandoc (required for rendering Jupyter notebooks):

   • Ubuntu/Debian: sudo apt install pandoc
   • macOS: brew install pandoc
   • Windows: choco install pandoc or download from https://pandoc.org/installing.html

3. Build the documentation:

   make html
   

The generated documentation will be in _build/html/.

Author

Ion Gabriel Ion, e-mail: ion.ion.gabriel@gmail.com