trainsum: Python QTT Arithmetic

• trainsum is a Python package for computations with multidimensional quantics tensor trains (QTTs) that enables efficient tensor arithmetic independent of shape or dimensionality.
• It employs the Array API standard and opt_einsum to perform addition, Einstein summations, and elementwise transformations using decomposition, variational, and cross interpolation methods.
• Key applications include simulation, data compression, machine learning, and spectral analysis, with implementations ranging from PDE solvers to Fourier transforms.

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 (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

The package also formalizes quantized grids for sampled functions. For a uniform discretization on $x\in[a,b]$ with $A(i) \equiv A(i_1,\dotsc,i_n),$0 points,

$A(i) \equiv A(i_1,\dotsc,i_n),$1

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 (Haubenwallner et al., 23 Feb 2026).

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 $A(i) \equiv A(i_1,\dotsc,i_n),$2, whereas trainsum is built for arbitrary factorization patterns and multidimensional layouts (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

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(i) \equiv A(i_1,\dotsc,i_n),$3
• 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 (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

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) \equiv A(i_1,\dotsc,i_n),$4 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 (Haubenwallner et al., 23 Feb 2026, Calvi et al., 2017).

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 (Haubenwallner et al., 23 Feb 2026).

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 $A(i) \equiv A(i_1,\dotsc,i_n),$5 against an exact operation $A(i) \equiv A(i_1,\dotsc,i_n),$6. In the notation of the paper,

$A(i) \equiv A(i_1,\dotsc,i_n),$7

and the local optimization step for a normalization-center core $A(i) \equiv A(i_1,\dotsc,i_n),$8 is derived from

$A(i) \equiv A(i_1,\dotsc,i_n),$9

Multi-core super-core updates are also supported (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

Several structured QTT constructions are given analytically. Exponentials factorize into a rank-1 tensor train: $$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,$$0 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 $$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,$$1. The package also provides QTT forms for shift matrices, Toeplitz tensors, and a discrete Fourier transform (Haubenwallner et al., 23 Feb 2026).

The Fourier constructor is the clearest numerical benchmark reported for the package. For a size $$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,$$2 Fourier matrix, the approximation distances to the exact transform are

$$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,$$3

for maximum ranks

$$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,$$4

This shows a pronounced rank–accuracy tradeoff: increasing the allowed TT rank sharply improves approximation quality (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

For eigenvalue problems $$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,$$5, 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 $$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,$$6, 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 (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

The broader computational significance lies in the quantics viewpoint. By factorizing a large dimension into many small ones, a problem of size $$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,$$7 can often be represented with complexity growing roughly polylogarithmically in $$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,$$8 when low QTT ranks exist. This makes the package most useful when the target tensor or function exhibits exploitable low-rank structure after quantization (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

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 (Haubenwallner et al., 23 Feb 2026).

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 (Calvi et al., 2017). 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 (Haubenwallner et al., 23 Feb 2026).