torchtt package

Submodules

torchtt.cpp module

Module for the C++ backend.

torchtt.cpp.cpp_avaible()[source]

Return True if C++ backend is available.

Returns:

True if C++ backend is available and False otherwise.

Return type:

bool

torchtt.errors module

Contains the errors used in the torchtt package.

exception torchtt.errors.IncompatibleTypes[source]

Bases: Exception

The function arguments are not compatible.

Usually means that a TT matrix was passed as argument instead of a TT tensor (or viceversa).

exception torchtt.errors.InvalidArguments[source]

Bases: Exception

The arguments are not valid.

The arguments passed are not of valid type.

exception torchtt.errors.RankMismatch[source]

Bases: Exception

The TT-ranks do not match.

This means that the inputs shapes that do not match.

exception torchtt.errors.ShapeMismatch[source]

Bases: Exception

The shape of the tensors does not match.

This means that the inputs have shapes that do not match.

torchtt.grad module

Adds AD functionality to torchtt.

torchtt.grad.grad(val, tens, core_indices=None)[source]

Compute the gradient w.r.t. the cores of the given TT-tensor (or TT-matrix).

Parameters:
  • val (torch.tensor) – Scalar tensor that has to be differentiated.

  • tens (torchtt.TT) – The given tensor.

  • core_indices (list[int], optional) – The list of cores to construct the gradient. If None is provided, all the cores are watched. Defaults to None.

Returns:

the list of cores representing the derivative of the expression w.r.t the tensor.

Return type:

list[torch.tensor]

torchtt.grad.grad_list(val, tensors, all_in_one=True)[source]

Compute the gradient w.r.t. the cores of several given TT-tensors (or TT-oeprators). Watch must be called on all of them beforehand.

Parameters:
  • val (torch.tensor) – scalar tensor to be differentiated.

  • tensors (list[torch.TT]) – the tensors with respect to which the differentiation is made.

  • all_in_one (bool, optional) – Put all the cores in one list or create a list of lists with the cores. Defaults to True.

Returns:

the resulting derivatives.

Return type:

list[list[torchtt.TT]]

torchtt.grad.unwatch(tens)[source]

Cancel the autograd graph recording.

Parameters:

tens (torchtt.TT) – the tensor.

torchtt.grad.watch(tens, core_indices=None)[source]

Watch the TT-cores of a given tensor. Necessary for autograd.

Parameters:
  • tens (torchtt.TT) – the TT-object to be watched.

  • core_indices (list[int], optional) – The list of cores to be watched. If None is provided, all the cores are watched. Defaults to None.

torchtt.grad.watch_list(tensors)[source]

Watch the TT-cores for amultiple tensors givgen in a list. Necessary for autograd.

Parameters:

tensors (list[torchtt.TT]) – the list of tensors to be wtched for autograd.

torchtt.interpolate module

Implements the cross approximation methods (DMRG).

class torchtt.interpolate.AmenCrossCallbacks(function, eval_mv, x, N, dtype, device)[source]

Bases: AmenCallbacks

compute_enrichment(k, state_dict, x_cores, z_cores, u, v)[source]
compute_x_bck(k, state_dict, x_cores, z_cores)[source]
compute_x_fwd(k, state_dict, x_cores, z_cores)[source]
compute_z_bck(k, state_dict, x_cores, z_cores)[source]
compute_z_fwd(k, state_dict, x_cores, z_cores, u, v)[source]
update_phis_bck(k, state_dict, x_cores, z_cores, swp, last)[source]
update_phis_fwd(k, state_dict, x_cores, z_cores, swp, last)[source]
torchtt.interpolate.dmrg_cross(function, N, eps=1e-09, nswp=10, x_start=None, kick=2, dtype=torch.float64, device=None, eval_vect=True, rmax=9223372036854775807, verbose=False, callback=None)[source]

Approximate a tensor in the TT format given that the individual entries are given using a function. The function is given as a function handle taking as arguments a matrix of integer indices.

Example

func = lambda I: 1/(2+I[:,0]+I[:,1]+I[:,2]+I[:,3]).to(dtype=torch.float64)
N = [20]*4
x = torchtt.interpolate.dmrg_cross(func, N, eps = 1e-7)
Parameters:
  • function (Callable) – function handle.

  • N (list[int]) – the shape of the tensor.

  • eps (float, optional) – the relative accuracy. Defaults to 1e-9.

  • nswp (int, optional) – number of iterations. Defaults to 20.

  • x_start (torchtt.TT, optional) – initial approximation of the output tensor (None coresponds to random initialization). Defaults to None.

  • kick (int, optional) – enrichment rank. Defaults to 2.

  • dtype (torch.dtype, optional) – the dtype of the result. Defaults to tn.float64.

  • device (torch.device, optional) – the device where the approximation will be stored. Defaults to None.

  • eval_vect (bool, optional) – not yet implemented. Defaults to True.

  • rmax (int, optional) – the maximum rank. Defaults to the maximum possible integer.

  • verbose (bool, optional) – display debug information to the console. Defaults to False.

  • callback (Callable, optional) – optional hook invoked at the end of every sweep as callback(tt, sweep, error), where tt is the current approximation (torchtt.TT), sweep is the 0-based sweep index (int) and error is the convergence metric for that sweep (float). If it returns False the sweeping is stopped early; any other return value continues. Useful for logging or custom stopping criteria. Defaults to None.

Returns:

the result.

Return type:

torchtt.TT

torchtt.interpolate.function_interpolate(function, x, eps=1e-09, start_tens=None, nswp=20, kick=2, kick2=0, dtype=torch.float64, rmax=9223372036854775807, method='dmrg', verbose=False, callback=None)[source]

Interpolate a function using tensor train cross approximation.

Parameters:
  • function (Callable) – Function to interpolate.

  • x (torchtt.TT or list[torchtt.TT]) – The points at which to evaluate the function.

  • eps (float, optional) – The desired relative error. Defaults to 1e-9.

  • start_tens (torchtt.TT, optional) – Initial tensor train approximation. Defaults to None.

  • nswp (int, optional) – Number of sweeps. Defaults to 20.

  • kick (int, optional) – Rank enrichment. Defaults to 2.

  • kick2 (int, optional) – Secondary rank enrichment (meant for amen method). Defaults to 0.

  • dtype (torch.dtype, optional) – The datatype of the result. Defaults to tn.float64.

  • rmax (int, optional) – Maximum allowed rank. Defaults to sys.maxsize.

  • method (str, optional) – Method to use (‘dmrg’ or ‘amen’). Defaults to ‘dmrg’.

  • verbose (bool, optional) – If True, display information. Defaults to False.

  • callback (Callable, optional) – optional hook invoked at the end of every sweep as callback(tt, sweep, error), where tt is the current approximation (torchtt.TT), sweep is the 0-based sweep index (int) and error is the convergence metric for that sweep (float). If it returns False the sweeping is stopped early; any other return value continues. Useful for logging or custom stopping criteria. Defaults to None.

Raises:

ValueError – If the method is not ‘dmrg’ or ‘amen’.

Returns:

The interpolated tensor.

Return type:

torchtt.TT

torchtt.manifold module

Manifold gradient module.

torchtt.manifold.riemannian_gradient(x, func)[source]

Compute the Riemannian gradient using AD.

Parameters:
  • x (torchtt.TT) – the point on the manifold where the gradient is computed.

  • func ([type]) – function that has to be differentiated. The function takes as only argument torchtt.TT instances.

Returns:

the gradient projected on the tangent space of x.

Return type:

torchtt.TT

torchtt.manifold.riemannian_projection(Xspace, z)[source]

Project the tensor z onto the tangent space defined at xspace

Parameters:
  • Xspace (torchtt.TT) – the target where the tensor should be projected.

  • z (torchtt.TT) – the tensor that should be projected.

Raises:

IncompatibleTypes – Both must be of same type.

Returns:

the projection.

Return type:

torchtt.TT

torchtt.nn module

Implements a basic TT layer for constructing deep TT networks.

class torchtt.nn.AffineTransform(dim)[source]

Bases: Transform

Affine transformation block: z = R(theta) diag(exp(a)) x + b

Where R(theta) is a rotation matrix parameterized by Givens angles. det J = prod(exp(a))

forward(x, params)[source]

Applies the transformation. :param x: Input tensor of shape (…, dim). :type x: torch.Tensor :param params: Parameter tensor of shape (…, input_requirement). :type params: torch.Tensor

Returns:

Transformed tensor. det_jac (torch.Tensor or float): Determinant of the Jacobian.

Return type:

z (torch.Tensor)

input_requirement()[source]

Returns the number of parameters required from the NN per sample.

class torchtt.nn.ComposedTransform(transforms)[source]

Bases: Transform

Chains multiple transformations.

forward(x, params)[source]

Applies the transformation. :param x: Input tensor of shape (…, dim). :type x: torch.Tensor :param params: Parameter tensor of shape (…, input_requirement). :type params: torch.Tensor

Returns:

Transformed tensor. det_jac (torch.Tensor or float): Determinant of the Jacobian.

Return type:

z (torch.Tensor)

input_requirement()[source]

Returns the number of parameters required from the NN per sample.

class torchtt.nn.CompressedTTLayer(N_in, N_out, R_layer, R_output, activation=<built-in method relu of type object>, bias=True, dtype=torch.float32)[source]

Bases: Module

A Tensor Train (TT) neural network layer that operates directly on TT objects and applies nonlinear activation between TT cores during multiplication.

This layer is inspired by Nonlinear Tensor Train formats for deep neural networks. Instead of computing the full dense tensor, it performs a fast matrix-vector-like multiplication of the layer’s TTM weights with the input TT object. The layer natively applies a nonlinear activation and an optional bias to the intermediate core representations during the contraction sweep. The intermediate representations are orthogonalized and truncated to maintain the compression, ensuring the output TT object’s ranks are strictly bounded by R_output.

forward(x)[source]

Forward pass for the CompressedTTLayer.

Computes the operation by multiplying the input TT with the layer’s TTM using a right-to-left sweep. During this sweep, the intermediate bias is added and the activation is applied to the core before the core is orthogonalized and truncated (using SVD) to strictly enforce the per-bond rank limits given by R_output.

Parameters:

x (torchtt.TT) – The input Tensor Train object.

Returns:

The output Tensor Train object representing the nonlinearly compressed forward pass, with ranks bounded by R_output.

Return type:

torchtt.TT

class torchtt.nn.LinearLayerTT(size_in, size_out, rank, dtype=torch.float32, initializer='He')[source]

Bases: Module

Basic class for TT layers. See Tensorizing Neural Networks for a detailed description. It can be used similarily to any layer from torch.nn. The output of the layer is \(\mathcal{LTT}(\mathsf{x}) =\mathsf{Wx}+\mathsf{b}\), where the tensor operator \(\mathsf{W}\) is represented in the TT format (with a fixed prescribed rank).

forward(x)[source]

Computes the output of the layer for the given input.

Supports trailing dimensiond broadcasting. If the input of the layer is set to [M1,...,Md] and a tensor od shape [...,M1,...,Md] is provided then the multiplication is performed along the last d dimensions.

Parameters:

x (torch.tensor) – input of the layer.

Returns:

output of the layer.

Return type:

torch.tensor

class torchtt.nn.Rank1Shear(dim, degree=2, clip=None)[source]

Bases: Transform

Volume-preserving rank-1 nonlinear shear: z = x + u * g(v^T x + c)

To ensure volume preservation (det J = 1), u is projected orthogonally to v such that v^T u = 0. The scalar function g(t) is a polynomial of degree D with no constant term: g(t) = sum_{p=1}^D alpha_p t^p. If clipping is applied, the displacement is squashed to prevent numerical overflow: g_clip(t) = clip * tanh(g(t) / clip).

forward(x, params)[source]

Applies the transformation. :param x: Input tensor of shape (…, dim). :type x: torch.Tensor :param params: Parameter tensor of shape (…, input_requirement). :type params: torch.Tensor

Returns:

Transformed tensor. det_jac (torch.Tensor or float): Determinant of the Jacobian.

Return type:

z (torch.Tensor)

input_requirement()[source]

Returns the number of parameters required from the NN per sample.

class torchtt.nn.SinhArcsinhWarp(dim)[source]

Bases: Transform

Elementwise sinh-arcsinh warp: z_i = sinh(e^{s_i} asinh(x_i) + b_i)

forward(x, params)[source]

Applies the transformation. :param x: Input tensor of shape (…, dim). :type x: torch.Tensor :param params: Parameter tensor of shape (…, input_requirement). :type params: torch.Tensor

Returns:

Transformed tensor. det_jac (torch.Tensor or float): Determinant of the Jacobian.

Return type:

z (torch.Tensor)

input_requirement()[source]

Returns the number of parameters required from the NN per sample.

class torchtt.nn.TTDensityLayer(N, R, basis, transform=None, dtype=torch.float32)[source]

Bases: Module

A TT Density Layer evaluating \(p(x) = p_{\mathrm{ref}}(T(x)) \, |\det J_T(x)|\).

forward(tts, x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_tt(tts)[source]
input_requireemnt()
input_requirement()[source]
property integration_weights
marginalize(tts, x, keep_indices)[source]
class torchtt.nn.Transform(dim)[source]

Bases: Module

Base class for diffeomorphism transformations used in TTDensityLayer.

forward(x, params)[source]

Applies the transformation. :param x: Input tensor of shape (…, dim). :type x: torch.Tensor :param params: Parameter tensor of shape (…, input_requirement). :type params: torch.Tensor

Returns:

Transformed tensor. det_jac (torch.Tensor or float): Determinant of the Jacobian.

Return type:

z (torch.Tensor)

input_requirement()[source]

Returns the number of parameters required from the NN per sample.

class torchtt.nn.TriangularPolyShear(dim, degree)[source]

Bases: Transform

Knothe-Rosenblatt-style triangular polynomial shear: z_i = x_i + sum_{j>i} q_{ij}(x_j)

Each q_{ij}(t) is a polynomial of degree D_poly with no constant term: q_{ij}(t) = sum_{p=1}^{D_poly} beta_{ij,p} t^p. det J = 1.

forward(x, params)[source]

Applies the transformation. :param x: Input tensor of shape (…, dim). :type x: torch.Tensor :param params: Parameter tensor of shape (…, input_requirement). :type params: torch.Tensor

Returns:

Transformed tensor. det_jac (torch.Tensor or float): Determinant of the Jacobian.

Return type:

z (torch.Tensor)

input_requirement()[source]

Returns the number of parameters required from the NN per sample.

class torchtt.nn.TriangularShear(dim)[source]

Bases: Transform

Unit upper-triangular linear shear (completes the affine family). det J = 1.

forward(x, params)[source]

Applies the transformation. :param x: Input tensor of shape (…, dim). :type x: torch.Tensor :param params: Parameter tensor of shape (…, input_requirement). :type params: torch.Tensor

Returns:

Transformed tensor. det_jac (torch.Tensor or float): Determinant of the Jacobian.

Return type:

z (torch.Tensor)

input_requirement()[source]

Returns the number of parameters required from the NN per sample.

torchtt.solvers module

System solvers in the TT format.

torchtt.solvers.amen_solve(A, b, nswp=22, x0=None, eps=1e-10, rmax=32768, max_full=256, kickrank=4, kick2=0, trunc_norm='res', local_solver=1, local_iterations=40, resets=2, verbose=False, preconditioner=None, use_cpp=True, band_diagonal=-1, use_single_precision=False)[source]

Solve a multilinear system \(\mathsf{Ax} = \mathsf{b}\) in the Tensor Train format.

This method implements the algorithm from Sergey V Dolgov, Dmitry V Savostyanov, Alternating minimal energy methods for linear systems in higher dimensions.

Example

import torchtt
A = torchtt.random([(4,4),(5,5),(6,6)],[1,2,3,1]) # create random matrix
x = torchtt.random([4,5,6],[1,2,3,1]) # invent a random solution
b = A @ x # compute the rhs
xx = torchtt.solvers.amen_solve(A,b) # solve
print((xx-x).norm()/x.norm()) # error
Parameters:
  • A (torchtt.TT) – the system matrix in TT.

  • b (torchtt.TT) – the right hand side in TT.

  • nswp (int, optional) – number of sweeps. Defaults to 22.

  • x0 (torchtt.TT, optional) – initial guess. In None is provided the initial guess is a ones tensor. Defaults to None.

  • eps (float, optional) – relative residual. Defaults to 1e-10.

  • rmax (int, optional) – maximum rank. Defaults to 100000.

  • max_full (int, optional) – local systems strictly smaller than this are solved directly, larger ones with the iterative local solver. Defaults to 256.

  • kickrank (int, optional) – rank enrichment. Defaults to 4.

  • kick2 (int, optional) – [description]. Defaults to 0.

  • trunc_norm (str, optional) – [description]. Defaults to ‘res’.

  • local_solver (int, optional) – choose local iterative solver: 1 for GMRES and 2 for BiCGSTAB. Defaults to 1.

  • local_iterations (int, optional) – number of GMRES iterations for the local subproblems. Defaults to 40.

  • resets (int, optional) – number of resets in the GMRES. Defaults to 2.

  • verbose (bool, optional) – choose whether to display or not additional information during the runtime. Defaults to True.

  • preconditioner (string, optional) – Choose the preconditioner for the local system. Possible values are None, ‘c’ (central Jacobi preconditioner). No preconditioner is used if None is provided. Defaults to None.

  • use_cpp (bool, optional) – use the C++ implementation of AMEn. Defaults to True.

  • band_diagonal (int, optional) – the TT cores of the matrix habe band diagonal structure. Defaults to -1.

Raises:
Returns:

the approximation of the solution in TT format.

Return type:

torchtt.TT

torchtt.solvers.cpp_enabled()[source]

Is the C++ backend enabled?

Returns:

the flag

Return type:

bool

Module contents

Provides Tensor-Train (TT) decomposition using pytorch as backend.

Contains routines for computing the TT decomposition and all the basisc linear algebra in the TT format. Additionally, GPU support can be used thanks to the pytorch backend. It also has linear solvers in TT and cross approximation as well as automatic differentiation.

class torchtt.TT(source, shape=None, eps=1e-10, rmax=9223372036854775807)[source]

Bases: object

property M

Return the “row” shape in case of TT matrices.

Raises:

IncompatibleTypes – The field is_ttm is defined only for TT matrices.

Returns:

the shape.

Return type:

list[int]

property N

Return the shape of a tensor or the “column” shape of a TT operator.

Returns:

the shape.

Return type:

list[int]

property R

The rank of the TT decomposition. It’s length should be len(R)==len(N)+1.

Returns:

the rank.

Return type:

list[int]

apply_mask(indices)[source]

Evaluate the tensor on the given index list.

Examples

x = torchtt.random([10,12,14],[1,4,5,1])
indices = torch.tensor([[0,0,0],[1,2,3],[1,1,1]])
val = x.apply_mask(indices)
Parameters:

indices (list[list[int]]) – the index list where the tensor should be evaluated. Length is M.

Returns:

the values of the tensor

Return type:

torch.tensor

clone()[source]

Clones the torchtt.TT instance. Similar to torch.tensor.clone().

Returns:

the cloned TT object.

Return type:

torchtt.TT

conj()[source]

Return the complex conjugate of a tensor in TT format.

Returns:

the complex conjugated tensor.

Return type:

torchtt.TT

cpu()[source]

Retrive the cores from the GPU.

Returns:

The TT-object on CPU.

Return type:

torchtt.TT

cuda(device=None)[source]

Return a torchtt.TT object on the CUDA device by cloning all the cores on the GPU.

Parameters:

device (torch.device, optional) – The CUDA device (None for CPU). Defaults to None.

Returns:

The TT-object. The TT-cores are on CUDA.

Return type:

torchtt.TT

detach()[source]

Detaches the TT tensor. Similar to torch.tensor.detach().

Returns:

the detached tensor.

Return type:

torchtt.TT

fast_matvec(other, eps=1e-12, initial=None, nswp=20, verb=False, use_cpp=True)[source]

Fast matrix vector multiplication A@x using DMRG iterations. Faster than traditional matvec + rounding.

Parameters:
  • other (torchtt.TT) – the TT tensor.

  • eps (float, optional) – relative accuracy for DMRG. Defaults to 1e-12.

  • initial (None|torchtt.TT, optional) – an approximation of the product (None means random initial guess). Defaults to None.

  • nswp (int, optional) – number of DMRG iterations. Defaults to 40.

  • verb (bool, optional) – show info for debug. Defaults to False.

  • use_cpp (bool, optional) – use the C++ implementation if available. Defaults to True.

Raises:
Returns:

the result.

Return type:

torchtt.TT

full()[source]

Return the full tensor. In case of a TTM, the result has the shape M1 x M2 x ... x Md x N1 x N2 x ... x Nd.

Returns:

the full tensor.

Return type:

torch.tensor

is_cuda()[source]

Return True if the tensor is on GPU.

Returns:

Is the torchtt.TT on GPU or not.

Return type:

bool

property is_ttm

Check whether the instance is a TT operator or not.

Returns:

the flag.

Return type:

bool

mprod(factor_matrices, mode)[source]

n-mode product.

Parameters:
  • factor_matrices (torch.tensor or list[torch.tensor]) – either a single matrix is directly provided or a list of matrices for product along multiple modes.

  • mode (int or list[int]) – the mode for the product. If factor_matrices is a torch.tensor then mode is an integer and the multiplication will be performed along a single mode. If factor_matrices is a list, the mode has to be list[int] of equal size.

Raises:
Returns:

the result

Return type:

torchtt.TT

norm(squared=False)[source]

Computes the frobenius norm of a TT object.

Parameters:

squared (bool, optional) – returns the square of the norm if True. Defaults to False.

Returns:

the norm.

Return type:

torch.tensor

numpy()[source]

Return the full tensor as a numpy.array. In case of a TTM, the result has the shape M1 x M2 x ... x Md x N1 x N2 x ... x Nd. If it is involved in an AD graph, an error will occur.

Returns:

the full tensor in numpy.

Return type:

numpy.array

qtt_to_tens(original_shape)[source]

Transform a tensor back from QTT.

Parameters:

original_shape (list) – the original shape.

Raises:
Returns:

the folded tensor.

Return type:

torchtt.TT

reduce_dims(exclude=[])[source]

Reduces the size 1 modes of the TT-object. At least one mode should be larger than 1.

Parameters:

exclude (list, optional) – Indices to exclude. Defaults to [].

round(eps=1e-12, rmax=9223372036854775807)[source]

Implements the rounding operations within a given tolerance epsilon. The maximum rank is also provided.

Parameters:
  • eps (float, optional) – the relative accuracy. Defaults to 1e-12.

  • rmax (int, optional) – the maximum rank. Defaults to the maximum possible integer.

Returns:

the result.

Return type:

torchtt.TT

set_core(k, core)[source]

Replaces the k-th TT core. This can change the mode size of the TT object.

Parameters:
  • k (int) – the core index.

  • core (torch.tensor) – _description_

Raises:
sum(index=None)[source]

Contracts a tensor in the TT format along the given indices and retuyrns the resulting tensor in the TT format. If no index list is given, the sum over all indices is performed.

Examples


a = torchtt.ones([3,4,5,6,7]) print(a.sum()) print(a.sum([0,2,4])) print(a.sum([1,2])) print(a.sum([0,1,2,3,4]))

Parameters:

index (int | list[int] | None, optional) – the indices along which the summation is performed. None selects all of them. Defaults to None.

Raises:

InvalidArguments – Invalid index.

Returns:

the result.

Return type:

torchtt.TT/torch.tensor

t()[source]

Returns the transpose of a given TT matrix.

Returns:

the transpose.

Return type:

torchtt.TT

Raises:

InvalidArguments – Has to be TT matrix.

to(device=None, dtype=None)[source]

Moves the TT instance to the given device with the given dtype.

Parameters:
  • device (torch.device, optional) – The desired device. If none is provided, the device is the CPU. Defaults to None.

  • dtype (torch.dtype, optional) – The desired dtype (torch.float64, torch.float32,…). If None is provided the dtype is not changed. Defaults to None.

to_qtt(eps=1e-12, mode_size=2, rmax=9223372036854775807)[source]

Converts a tensor to the QTT format: N1 x N2 x … x Nd -> mode_size x mode_size x … x mode_size. The product of the mode sizes should be a power of mode_size. The tensor in QTT can be converted back using the qtt_to_tens() method.

Examples

x = torchtt.random([16,8,64,128],[1,2,10,12,1])
x_qtt = x.to_qtt()
print(x_qtt)
xf = x_qtt.qtt_to_tens(x.N) # a TT-rounding is recommended.
Parameters:
  • eps (float,optional) – the accuracy. Defaults to 1e-12.

  • mode_size (int, optional) – the size of the modes. Defaults to 2.

  • rmax (int) – the maximum rank. Defaults to the maximum possible integer.

Raises:
  • ShapeMismatch – Only quadratic TTM can be tranformed to QTT.

  • ShapeMismatch – Reshaping error: check if the dimensions are powers of the desired mode size.

Returns:

the resulting reshaped tensor.

Return type:

torchtt.TT

to_ttm()[source]

Converts a TT-tensor to the TT-matrix format. In the tensor has the shape N1 x … x Nd, the result has the shape N1 x … x Nd x 1 x … x 1.

Returns:

the result

Return type:

torch.TT

class torchtt.Timer(name='')[source]

Bases: object

torchtt.amen_mm(A, B, nswp=22, X0=None, eps=1e-10, rmax=1024, kickrank=4, kick2=0, verbose=False)[source]

Perform the TTM-TTM product using AMEn optimization. Suited when the operators have high ranks, but the result is expected to be low rank.

Parameters:
  • A (torchtt.TT) – the first TTM.

  • B (torchtt.TT) – the second TTM.

  • nswp (int, optional) – number of sweeps. Defaults to 22.

  • X0 (torchtt.TT, optional) – initial guess (None means no initial guess). Defaults to None.

  • eps (float, optional) – realtive tolerance. Defaults to 1e-10.

  • rmax (int, optional) – maximum rank. Defaults to 1024.

  • kickrank (int, optional) – kickrank. Defaults to 4.

  • kick2 (int, optional) – kick2. Defaults to 0.

  • verbose (bool, optional) – show debug info. Defaults to False.

Returns:

the result.

Return type:

torchtt.TT

torchtt.amen_mv(A, b, nswp=22, x0=None, eps=1e-10, rmax=1024, kickrank=4, kick2=0, verbose=False, use_cpp=True)[source]

Compute the matrix vector product between a TTM and a TT. Suited when the output is expected to be low rank.

Parameters:
  • A (torchtt.TT) – the matrix in TT.

  • b (torchtt.TT) – the tensor TT.

  • nswp (int, optional) – number of sweeps. Defaults to 22.

  • x0 (torchtt.TT, optional) – initial guess. In None is provided the initial guess is a ones tensor. Defaults to None.

  • eps (float, optional) – relative residual. Defaults to 1e-10.

  • rmax (int, optional) – maximum rank. Defaults to 100.

  • kickrank (int, optional) – rank enrichment. Defaults to 4.

  • kick2 (int, optional) – [description]. Defaults to 0.

  • verbose (bool, optional) – choose whether to display or not additional information during the runtime. Defaults to True.

  • use_cpp (bool, optional) – use the C++ implementation of AMEn. Defaults to True.

Raises:
Returns:

the approximation of the solution in TT format.

Return type:

torchtt.TT

torchtt.bilinear_form(x, A, y)[source]

Computes the bilinear form x^T A y for TT tensors:

Parameters:
Raises:
Returns:

the result of the bilienar form as tensor with 1 element.

Return type:

torch.tensor

torchtt.cat(tensors, dim=0)[source]

Concatenate tensors in the TT format along a given dimension dim. Only works for TT tensors and not TT matrices.

Examples

import torchtt
import torch


a1 = torchtt.randn((3,4,2,6,7), [1,2,3,4,2,1])
a2 = torchtt.randn((3,4,8,6,7), [1,3,1,7,5,1])
a3 = torchtt.randn((3,4,15,6,7), [1,3,10,2,4,1])

a = torchtt.cat((a1,a2,a3),2)

af = torch.cat((a1.full(), a2.full(),
print(torch.linalg.norm(a.full()-af))

`

Parameters:
  • tensors (tuple[TT]) – the tensors to be concatenated. Their mode sizes must match for all modex except the concatenating dimension.

  • dim (int, optional) – The dimension to be concatenated after. Defaults to 0.

Raises:
  • InvalidArguments – Not implemented for tensor matrices.

  • InvalidArguments – The mode sizes must be the same on the nonconcatenated dimensions for all the provided tensors.

  • InvalidArguments – The tensors must have the same number of dimensions.

Returns:

the result.

Return type:

torchtt.TT

torchtt.diag(input)[source]

Creates diagonal TT matrix from TT tensor or extracts the diagonal of a TT matrix:

  • If a TT matrix is provided the result is a TT tensor representing the diagonal :math:` mathsf{x}_{i_1…i_d} = mathsf{A}_{i_1…i_d,i_1…i_d} `

  • If a TT tensor is provided the result is a diagonal TT matrix with the entries :math:` mathsf{A}_{i_1…i_d,j_1…j_d} = mathsf{x}_{i_1…i_d} delta_{i_1}^{j_1} cdots delta_{i_d}^{j_d} `

Parameters:

input (TT) – the input.

Raises:

InvalidArguments – Input must be a torchtt.TT instance.

Returns:

the result.

Return type:

torchtt.TT

torchtt.dmrg_hadamard(x, y, z0=None, nswp=20, eps=1e-12, rmax=32768, kickrank=4, verb=False, use_cpp=True)[source]

Perform fast elementwise multiplication z = x * y in the TT using the DMRG algorithm. C++ backend not yet ready if available.

Parameters:
  • z (TT) – TT tensor

  • x (TT) – TT tensor

  • z0 (TT, optional) – initial guess of the result (if None is provided a random tensor is generated as a guess). Defaults to None.

  • nswp (int, optional) – numebr of sweeps. Defaults to 20.

  • eps (float, optional) – relative accuracy. Defaults to 1e-12.

  • rmax (int, optional) – maximum rank. Defaults to 32768.

  • kickrank (int, optional) – kickrank. Defaults to 4.

  • verb (bool, optional) – show debug info. Defaults to False.

  • use_cpp (bool, optional) – flag to choose between the python and C++ implementation (if available). Defaults to False.

Returns:

the result.

Return type:

TT

torchtt.dot(a, b, axis=None)[source]

Computes the dot product between 2 tensors in TT format. If both a and b have identical mode sizes the result is the dot product. If a and b have inequal mode sizes, the function perform index contraction. The number of dimensions of a must be greater or equal as b. The modes of the tensor a along which the index contraction with b is performed are given in axis. For the compelx case (a,b) = b^H . a.

Examples

a = torchtt.randn([3,4,5,6,7],[1,2,2,2,2,1])
b = torchtt.randn([3,4,5,6,7],[1,2,2,2,2,1])
c = torchtt.randn([3,5,6],[1,2,2,1])
print(torchtt.dot(a,b))
print(torchtt.dot(a,c,[0,2,3]))
Parameters:
  • a (torchtt.TT) – the first tensor.

  • b (torchtt.TT) – the second tensor.

  • axis (list[int], optional) – the mode indices for index contraction. Defaults to None.

Raises:
  • InvalidArguments – Both operands should be TT instances.

  • NotImplementedError – Operation not implemented for TT-matrices.

  • ShapeMismatch – Operands are not the same size.

  • ShapeMismatch – Number of the modes of the first tensor must be equal with the second.

Returns:

the result. If no axis index is provided the result is a scalar otherwise a torchtt.TT object.

Return type:

float or torchtt.TT

torchtt.elementwise_divide(x, y, eps=1e-12, starting_tensor=None, nswp=50, kick=4, local_iterations=40, resets=2, preconditioner=None, verbose=False)[source]

Perform the elemntwise division x/y of two tensors in the TT format using the AMEN method. Use this method if different AMEN arguments are needed. This method does not check the validity of the inputs.

Parameters:
  • x (torchtt.TT or scalar) – first tensor (can also be scalar of type float, int, torch.tensor with shape (1)).

  • y (torchtt.TT) – second tensor.

  • eps (float, optional) – relative acccuracy. Defaults to 1e-12.

  • starting_tensor (torchtt.TT or None, optional) – initial guess of the result (None for random initial guess). Defaults to None.

  • nswp (int, optional) – number of iterations. Defaults to 50.

  • kick (int, optional) – size of rank enrichment. Defaults to 4.

  • local_iterations (int, optional) – the number of iterations for the local iterative solver. Defaults to 40.

  • resets (int, optional) – the number of restarts in the GMRES solver. Defaults to 2.

  • preconditioner (string, optional) – Use preconditioner for the local solver (possible vaules None, ‘c’). Defaults to None.

  • verbose (bool, optional) – display debug info. Defaults to False.

Returns:

the result

Return type:

torchtt.TT

torchtt.eye(shape, dtype=torch.float64, device=None)[source]

Construct the TT decomposition of a multidimensional identity matrix. all the TT ranks are 1.

Parameters:
  • shape (list[int]) – the shape.

  • dtype (torch.dtype, optional) – the dtype of the returned tensor. Defaults to tn.float64.

  • device (torch.device, optional) – the device where the TT cores are created (None means CPU). Defaults to None.

Returns:

the one tensor.

Return type:

torchtt.TT

torchtt.fast_mm(tt_a, tt_b, eps=1e-10, rmax=None)[source]

Performs the matmat product between a TTM and a TTM. Equivalent to (tt_a * tt_b).round(eps). Method described in [https://arxiv.org/pdf/2410.19747](https://arxiv.org/pdf/2410.19747).

Parameters:
  • tt_a (torchtt.TT) – the first operand. Must be a TTM.

  • tt_b (torchtt.TT) – the second operand. Must be TTM.

  • eps (float, optional) – Relative tolerance. Defaults to 1e-10.

  • rmax (int, optional) – maximum rank. Defaults to None.

Raises:
Returns:

the result. This is a TTM.

Return type:

torchtt.TT

torchtt.fast_mv(tt_a, tt_b, eps=1e-10, rmax=None)[source]

Performs the matvec product between a TTM and a TT. Equivalent to (tt_a * tt_b).round(eps). Method described in [https://arxiv.org/pdf/2410.19747](https://arxiv.org/pdf/2410.19747).

Parameters:
  • tt_a (torchtt.TT) – the first operand. Must be a TTM.

  • tt_b (torchtt.TT) – the second operand. Must be TT.

  • eps (float, optional) – Relative tolerance. Defaults to 1e-10.

  • rmax (int, optional) – maximum rank. Defaults to None.

Raises:
  • InvalidArguments – The first should be e TTM and the second a TT.

  • ShapeMismatch – The shapes of the two operands must be compatible: tt_a.N == tt_b.N.

Returns:

the result. This is a TT.

Return type:

torchtt.TT

torchtt.kron(first, second)[source]

Computes the tensor Kronecker product. If None is provided as input the reult is the other tensor. If A is N_1 x … x N_d and B is M_1 x … x M_p, then kron(A,B) is N_1 x … x N_d x M_1 x … x M_p

Parameters:
Raises:
Returns:

the result.

Return type:

torchtt.TT

torchtt.load(path)[source]

Load a torchtt.TT object from a file.

Examples

import torchtt
#generate a TT object
A = torchtt.randn([10,20,30,40,4,5],[1,6,5,4,3,2,1])
# save the TT object
torchtt.save(A,"./test.TT")
# load the TT object
B = torchtt.load("./test.TT")
# the loaded should be the same
print((A-B).norm()/A.norm())
Parameters:

path (str) – the file name.

Returns:

the tensor.

Return type:

torchtt.TT

torchtt.meshgrid(vectors)[source]

Creates a meshgrid of torchtt.TT objects. Similar to numpy.meshgrid or torch.meshgrid. The input is a list of d torch.tensor vectors of sizes N_1, … ,N_d The result is a list of torchtt.TT instances of shapes N1 x … x Nd.

Parameters:

vectors (list[torch.tensor]) – the vectors (1d tensors).

Returns:

the resulting meshgrid.

Return type:

list[TT]

torchtt.numel(tensor)[source]

Return the number of entries needed to store the TT cores for the given tensor.

Parameters:

tensor (torchtt.TT) – the TT representation of the tensor.

Returns:

number of floats stored for the TT decomposition.

Return type:

int

torchtt.ones(shape, dtype=torch.float64, device=None)[source]

Construct a tensor that contains only ones. the shape can be a list of ints or a list of tuples of ints. The second case creates a TT matrix.

Parameters:
  • shape (list[int] or list[tuple[int]]) – the shape.

  • dtype (torch.dtype, optional) – the dtype of the returned tensor. Defaults to tn.float64.

  • device (torch.device, optional) – the device where the TT cores are created (None means CPU). Defaults to None.

Raises:

InvalidArguments – Shape must be a list.

Returns:

the one tensor.

Return type:

torchtt.TT

torchtt.pad(tensor, padding, value=0.0)[source]

Pad a tensor in the TT format. The padding argument is a tuple of tuples ((b1, a1), (b2, a2), … , (bd, ad)). Each dimension is padded with bk at the beginning and ak at the end. The padding value is constant and is given as the argument value. In case of a TT operator, duiagual padding is performed. On the diagonal, the provided value is inserted.

Parameters:
  • tensor (TT) – the tensor to be padded.

  • padding (tuple(tuple(int))) – the paddings.

  • value (float, optional) – the value to pad. Defaults to 0.0.

Raises:

InvalidArguments – The number of paddings should not exceed the number of dimensions of the tensor.

Returns:

the result.

Return type:

TT

torchtt.permute(input, dims, eps=1e-12)[source]

Permutes the dimensions of the tensor. Works similarily to torch.permute. Works like a bubble sort for both TT tensors and TT matrices.

Examples

x_tt = torchtt.random([5,6,7,8,9],[1,2,3,4,2,1])
xp_tt = torchtt.permute(x_tt, [4,3,2,1,0], 1e-10)
print(xp_tt) # the shape of this tensor should be [9,8,7,6,5]
Parameters:
  • input (torchtt.TT) – the input tensor.

  • dims (list[int]) – the order of the indices in the new tensor.

  • eps (float, optional) – the relative accuracy of the decomposition. Defaults to 1e-12.

Raises:
Returns:

the resulting tensor.

Return type:

torchtt.TT

torchtt.randn(N, R, var=1.0, dtype=torch.float64, device=None)[source]

A torchtt.TT tensor of shape N = [N1 x … x Nd] and rank R is returned. The entries of the fuill tensor are alomst normal distributed with the variance var.

Parameters:
  • N (list[int]) – the shape.

  • R (list[int]) – the rank.

  • var (float, optional) – the variance. Defaults to 1.0.

  • dtype (torch.dtype, optional) – the dtype of the returned tensor. Defaults to tn.float64.

  • device (torch.device, optional) – the device where the TT cores are created (None means CPU). Defaults to None.

Returns:

the result.

Return type:

torchtt.TT

torchtt.random(N, R, dtype=torch.float64, device=None)[source]

Returns a tensor of shape N with random cores of rank R. Each core is a normal distributed with mean 0 and variance 1. Check also the method torchtt.randn()for better random tensors in the TT format.

Parameters:
  • N (list[int] or list[tuple[int]]) – the shape of the tensor. If the elements are tuples of integers, we deal with a TT-matrix.

  • R (list[int] or int) – can be a list if the exact rank is specified or an integer if the maximum rank is secified.

  • dtype (torch.dtype, optional) – the dtype of the returned tensor. Defaults to tn.float64.

  • device (torch.device, optional) – the device where the TT cores are created (None means CPU). Defaults to None.

Raises:

InvalidArguments – Check if N and R are right.

Returns:

the result.

Return type:

torchtt.TT

torchtt.rank1TT(elements)[source]

Compute the rank 1 TT from a list of vectors (or matrices).

Parameters:

elements (list[torch.tensor]) – the list of vectors (or matrices in case a TT matrix should be created).

Returns:

the resulting TT object.

Return type:

torchtt.TT

torchtt.reshape(tens, shape, eps=1e-16, rmax=9223372036854775807)[source]

Reshapes a torchtt.TT tensor in the TT format. A rounding is also performed.

Parameters:
  • tens (torchtt.TT) – the input tensor.

  • shape (list[int] or list[tuple[int]]) – the desired shape. In the case of a TT operator the shape has to be given as list of tuples of ints [(M1,N1),…,(Md,Nd)].

  • eps (float, optional) – relative accuracy. Defaults to 1e-16.

  • rmax (int, optional) – maximum rank. Defaults to the maximum possible integer.

Raises:

ShapeMismatch – The product of modes should remain equal. Check the given shape.

Returns:

the resulting tensor.

Return type:

torchtt.TT

torchtt.save(tensor, path)[source]

Save a torchtt.TT object in a file.

Examples

import torchtt
#generate a TT object
A = torchtt.randn([10,20,30,40,4,5],[1,6,5,4,3,2,1])
# save the TT object
torchtt.save(A,"./test.TT")
# load the TT object
B = torchtt.load("./test.TT")
# the loaded should be the same
print((A-B).norm()/A.norm())
Parameters:
  • tensor (torchtt.TT) – the tensor to be saved.

  • path (str) – the file name.

Raises:

InvalidArguments – First argument must be a torchtt.TT instance.

torchtt.shape_mn_to_tuple(M, N)[source]

Convert the shape of a TTM from row/column format to tuple format.

Parameters:
  • M (list[int]) – row shapes.

  • N (list[int]) – column shapes.

Returns:

shape.

Return type:

list[tuple[int]]

torchtt.shape_tuple_to_mn(shape)[source]

Convert the shape of a TTM from tuple format to row and column shapes.

Parameters:

shape (list[tuple[int]]) – shape.

Returns:

still the shape.

Return type:

tuple[list[int],list[int]]

torchtt.zeros(shape, dtype=torch.float64, device=None)[source]

Construct a tensor that contains only zeros. the shape can be a list of ints or a list of tuples of ints. The second case creates a TT matrix.

Parameters:
  • shape (list[int] | list[tuple[int]]) – the shape.

  • dtype (torch.dtype, optional) – the dtype of the returned tensor. Defaults to tn.float64.

  • device (torch.device, optional) – the device where the TT cores are created (None means CPU). Defaults to None.

Raises:

InvalidArguments – Shape must be a list.

Returns:

the zero tensor.

Return type:

torchtt.TT