tntorch – Tensor Network Learning with PyTorch¶
This is a PyTorch-powered library for tensor modeling and learning that features transparent support for the tensor train (TT) model, CANDECOMP/PARAFAC (CP), the Tucker model, and more. Supported operations (CPU and GPU) include:
- Basic and fancy indexing of tensors, broadcasting, assignment, etc.
- Tensor decomposition and reconstruction
- Element-wise and tensor-tensor arithmetics
- Building tensors from black-box functions using cross-approximation
- Statistics and sensitivity analysis
- Optimization using autodifferentiation, useful for e.g. regression or classification
- Misc. operations on tensors: stacking, unfolding, sampling, derivating, etc.
Get the Code¶
You can clone the project from tntorch’s GitHub page:
git clone https://github.com/rballester/tntorch.git
or get it as a zip file.
Installation¶
The main dependencies are NumPy and PyTorch (we recommend to install those with Conda or Miniconda). To install tntorch, run:
cd tntorch
pip install .
First Steps¶
Some basic tensor manipulation:
import tntorch as tn
t = tn.ones(64, 64) # 64 x 64 tensor, filled with ones
t = t[:, :, None] + 2*t[:, None, :] # Singleton dimensions, broadcasting, and arithmetics
print(tn.mean(t)) # Result: 3
Decomposing a tensor:
import tntorch as tn
data = ... # A NumPy or PyTorch tensor
t1 = tn.Tensor(data, ranks_cp=5) # A CP decomposition
t2 = tn.Tensor(data, ranks_tucker=5) # A Tucker decomposition
t3 = tn.Tensor(data, ranks_tt=5) # A tensor train decomposition
To get fully on board, check out the complete documentation: