---
title: 'trainsum: Python QTT Arithmetic'
url: https://www.emergentmind.com/topics/trainsum
type: topic
---

# trainsum: Python QTT Arithmetic

trainsum is a Python package for computations with **multidimensional quantics tensor trains (QTTs)**. It uses the **Array API standard** together with **`opt_einsum`** to approximate tensors or functions by tensor trains independent of their shape or dimensionality, and once approximated it can perform normal arithmetic operations with quantics tensor trains, including **addition**, **Einstein summations**, and **element-wise transformations**. The package is intended for generic computations with applications in **simulation**, **data compression**, **machine learning**, and **data analysis** [2602.20226].

## 1. Mathematical basis: quantization and tensor-train representation

The central idea of trainsum is to combine tensor-train compression with **quantics tensorization**. An original index \(i\) is factorized into several smaller indices,
\[
A(i) \equiv A(i_1,\dotsc,i_n),
\]
with
\[
i=\sum^n_{q=1} \left(\prod_{r=q+1}^n b^i_r\right) i_q \equiv \sum^n_{q=1} c^i_q i_q,
\qquad
c^i_q=\prod_{r=q+1}^n b^i_r,
\]
and
\[
\dim(i)=c_0=\prod_{q=1}^n b^i_q.
\]
This is the quantization step: a long axis is reshaped into multiple smaller “digit” axes. Unlike many QTT approaches that assume binary decomposition, trainsum allows **any factorization whose product matches the original size** [2602.20226].

After factorization, a tensor-train approximation is written as
\[
A(i)=A(i_1,\dotsc,i_n)=C_1(i_1)\cdot C_2(i_2)\cdots C_n(i_n),
\]
where each \(C_q(i_q)\) is a matrix of size \(d_q\times d_{q+1}\), with \(d_q\) the TT bond dimensions or ranks. For higher-order tensors the same construction is applied to grouped external indices, for example
\[
A(i,j,k)=C_1(i_1,j_1,k_1)\cdot C_2(i_2,j_2,k_2)\cdots C_n(i_n,j_n,k_n).
\]
The ranks control compression quality and computational cost: small ranks enable efficient storage and arithmetic, while large ranks reduce the advantage of the representation [2602.20226].

The package also formalizes quantized grids for sampled functions. For a uniform discretization on \(x\in[a,b]\) with \(\dim(i)\) points,
\[
x(i)=\frac{b-a}{\dim(i)-1}\,i+a.
\]
After factorization of the index, the grid coordinate can be written in digit form. This is the basis for direct QTT constructions of sampled functions such as **exponentials**, **trigonometric functions**, and **polynomials** [2602.20226].

## 2. Shape independence and software architecture

A defining property of trainsum is that it operates **independent of tensor shape or dimensionality**. The package is designed for generic arithmetic with quantized tensor representations of arbitrary-sized dimensions, including tensors, sampled data, and functions on grids. This is important because many existing tensor-network tools focus on matrix product states in physics settings or assume dimensions of size \(2^n\), whereas trainsum is built for arbitrary factorization patterns and multidimensional layouts [2602.20226].

This flexibility appears in the shape model. The package supports both **“block”** and **“interleaved”** arrangements of factorized digits across TT cores, and it distinguishes between the full tensor shape and the TT layout through **`TrainShape`**. Arithmetic therefore depends not only on nominal tensor dimensions, but also on how quantized digits are grouped into cores [2602.20226].

Its implementation is organized around a small set of core classes:

- **`TrainSum`**: public API namespace and factory  
- **`Dimension`**: factorized dimension as a sequence of digits  
- **`Digit`**: stores each factorized sub-dimension, including `base` and `factor`  
- **`Domain`**: interval \([a,b]\)  
- **`UniformGrid`**: mapping between indices and physical coordinates  
- **`TrainShape`**: tensor-train shape and core grouping  
- **`TensorTrain`**: TT cores and associated shape  
- **`LinearMap`**: tensor-train linear operators for eigensolvers and linear solvers  

Backend portability is provided through the **Python Array API standard** and `array_api_compat`. Three backends are supported out of the box: **NumPy**, **CuPy**, and **Torch**. The current design assumes **mutable N-dimensional arrays with statically defined shapes**, which excludes frameworks such as **JAX** or **Dask** [2602.20226].

Another architectural feature is the use of **context managers** to select the active execution strategy. The package provides contexts for **exact**, **decomposition**, **variational**, **cross**, and **evaluation** modes. The same high-level operation can therefore be executed with different approximation semantics without changing the calling syntax [2602.20226].

## 3. Arithmetic model and approximation algorithms

trainsum groups its operations into three broad classes: linear operations, nonlinear elementwise operations, and manipulation or utility operations. The package explicitly distinguishes algorithm families for these operations [2602.20226].

| Operation | Supported modes | Function |
|---|---|---|
| Einsum | exact, decomposition, variational, cross | linear contractions |
| Addition | exact, decomposition, variational, cross | TT sum |
| Element-wise functions | cross only | entrywise transforms |

The **exact addition** rule is the standard tensor-train direct sum:
\[
A(i_1)\cdots A(i_n) + B(i_1)\cdots B(i_n) =
\begin{pmatrix} A(i_1) & B(i_1)\end{pmatrix}
\begin{pmatrix} A(i_2) & 0 \\ 0 & B(i_2)\end{pmatrix}
\cdots
\begin{pmatrix} A(i_n) \\ B(i_n)\end{pmatrix}.
\]
This preserves the represented tensor exactly but increases ranks. The same block-structured direct-sum pattern appears in chain-based tensor-network summation more generally, where corresponding cores are combined in block form while physical dimensions remain fixed [2602.20226], [1711.00701].

**Exact Einstein summation** is the package’s flagship linear operation. Contracting an original dimension means contracting **all of its quantized digits**. Local contractions are performed core by core, but the resulting bond dimensions are flattened products of the input bond dimensions before truncation. The immediate consequence is rank growth; repeated exact operations therefore become impractical unless the result is recompressed [2602.20226].

To control rank growth, trainsum implements several approximation schemes. **Decomposition** or **zip-up** algorithms compute local exact super-cores and then compress them with approximate matrix decompositions such as SVD. The process advances through the train while restoring normalization as needed. **Variational** algorithms are DMRG-like rank-reduction procedures that optimize a target train \(C\) against an exact operation \(\hat O(A,B)\). In the notation of the paper,
\[
\left|C-\hat O(A,B)\right|=0,
\]
and the local optimization step for a normalization-center core \(C(i_q)\) is derived from
\[
\frac{\partial}{\partial C^*(i_q)}\langle C|C\rangle
=
\frac{\partial}{\partial C^*(i_q)}\langle C|\hat O(A,B)\rangle.
\]
Multi-core super-core updates are also supported [2602.20226].

For **nonlinear elementwise operations**, trainsum uses **cross interpolation**. This extends CUR-style ideas to tensor trains and works by sampling selected entries or one-dimensional slices rather than requiring the full tensor. In the package, cross interpolation is used both to construct TT approximations from callable functions and to apply arbitrary scalar transforms to existing TT objects [2602.20226].

## 4. Constructors, workflow, and structured operators

A typical trainsum workflow begins by defining factorized dimensions and grids. Passing an integer to **`ts.dimension`** triggers prime factorization; passing a list specifies the factorization explicitly. For interval-based problems, **`Domain`** and **`UniformGrid`** define the mapping between digit indices and physical coordinates. From there, **`TrainShape`** specifies the tensor-train layout, including block or interleaved digit grouping [2602.20226].

The package offers three construction modes. First, it includes **structured exact constructors** for common objects and operators: **`exp`**, **`sin`**, **`cos`**, **`polyval`**, **`full`**, **`shift`**, **`toeplitz`**, and **`qft`**. Second, it can approximate a dense tensor or sampled array through active **decomposition** or **variational** contexts. Third, it can approximate a callable function through **cross** interpolation, which is especially useful when the full tensor is too large to materialize. Direct construction from explicit TT cores is also supported [2602.20226].

Several structured QTT constructions are given analytically. Exponentials factorize into a **rank-1 tensor train**:
\[
e^{v(x-x_0)} \approx e^{v(a-x_0)}\prod_{q=1}^n e^{v c_q^x x_q}.
\]
Cosine is represented as a sum of two rank-1 trains and therefore admits **rank 2** after regauging to a real-valued form. Polynomials follow a constructive representation of rank \(p+1\). The package also provides QTT forms for **shift matrices**, **Toeplitz tensors**, and a **discrete Fourier transform** [2602.20226].

The Fourier constructor is the clearest numerical benchmark reported for the package. For a size \(2^{13}\) Fourier matrix, the approximation distances to the exact transform are
\[
[2.6\times 10^1,\ 3.5\times 10^{-1},\ 4.0\times 10^{-6},\ 9.3\times 10^{-11}]
\]
for maximum ranks
\[
[2,4,8,16].
\]
This shows a pronounced rank–accuracy tradeoff: increasing the allowed TT rank sharply improves approximation quality [2602.20226].

Once a tensor train has been constructed, trainsum overloads ordinary Python arithmetic. Addition, multiplication-like contractions, matrix–vector products, powers, division, absolute value, and arbitrary elementwise transforms can all be performed in compressed form under the currently active context. Outputs can remain as tensor trains, be materialized with **`to_tensor`**, evaluated at selected indices, or analyzed with **`min_max`** utilities [2602.20226].

## 5. Solvers and application domains

Beyond arithmetic, trainsum includes tensor-structured solver infrastructure. A **`LinearMap`** is defined by an Einstein summation expression together with TT operands and a designated shape. The operator must map the input shape back to the same shape, with input and output digits aligned on the same cores [2602.20226].

For eigenvalue problems \(H\Psi=E\Psi\), the package implements **DMRG-style** solvers. The paper describes a workflow that combines a local solver such as **Lanczos**, a decomposition method such as **`svdecomposition`**, and a sweeping strategy over one or more cores. For linear systems \(Ax=b\), it provides analogous methods, including a **DMRG-like energy minimization** approach and the **AMEn solver** with gradient enrichment. The package also mentions approximate TT minima and maxima algorithms [2602.20226].

The applications highlighted for trainsum span several domains. In **simulation**, the paper mentions PDEs and finite-difference systems such as the **heat equation**, as well as eigenvalue problems including the **hydrogen atom**. In **data compression**, it mentions **image compression**. In **machine learning**, it points to tensor-train neural network components, with an **MNIST** example. In **data analysis**, it emphasizes **Fourier transforms**, **spectral analysis**, and **convolutions** via Toeplitz tensors [2602.20226].

The broader computational significance lies in the quantics viewpoint. By factorizing a large dimension into many small ones, a problem of size \(N\) can often be represented with complexity growing roughly **polylogarithmically** in \(N\) when low QTT ranks exist. This makes the package most useful when the target tensor or function exhibits exploitable low-rank structure after quantization [2602.20226].

## 6. Limitations and relation to broader tensor-network research

The package has several explicit constraints. The Array-API abstraction currently requires **mutable arrays with statically defined shapes**, which excludes some backends. Efficient Einstein summation is not guaranteed for arbitrary tensor-train layouts; efficient execution requires **partial shape matching** on contracted factorized indices and that the result remain a **linear tensor network**. Otherwise computation may become expensive or require dense fallback [2602.20226].

A second limitation is intrinsic to the mathematics rather than the implementation. Exact tensor-train arithmetic causes **rank explosion**. In trainsum this appears most clearly in exact addition and exact Einstein summation, where ranks grow additively or multiplicatively before truncation. The package therefore depends heavily on decomposition, variational compression, or cross interpolation to keep computations tractable [2602.20226].

The paper also identifies missing or unstable functionality. **Arbitrary slicing and assignment operations** are still absent, and the numerical **stability of some algorithms, especially cross interpolation, could be improved**. These are presented as ongoing research directions rather than resolved issues [2602.20226].

Within tensor-network research more broadly, trainsum sits close to the representation-level summation viewpoint formalized in work on the **sum of tensor networks**, where tensors with equal physical-mode dimensions and isomorphic topology are combined by concatenating physical nodes and placing contraction nodes on the superdiagonal of block tensors [1711.00701]. This suggests that trainsum operationalizes a direct-sum principle for chain-like tensor networks inside a software system focused on multidimensional **QTT arithmetic**, rather than presenting tensor trains only as a storage format [2602.20226].

Source: https://www.emergentmind.com/topics/trainsum