Skip to content

Tensor Networks

tn

Tensor Network utilities and public classes.

CPNetwork dataclass 

CPNetwork(
    m_order: int | Sequence[int],
    rank: int,
    d_dim: int,
    seed: int | None = None,
    dtype: dtype = jnp.float32,
)

Bases: TensorNetwork

Canonical Polyadic (CP) tensor network representation.

Represents a tensor using the Canonical Polyadic (CP) decomposition, i.e., as a sum of rank-1 outer products. Each mode is represented by a factor matrix of shape (m_i, rank), where m_i is the local dimension of that mode.

Parameters:

  • m_order (int | Sequence[int]) –

    Local dimension per mode. If a single integer is provided, it is broadcast to all d_dim modes.

  • rank (int) –

    CP rank (number of rank-1 components).

  • d_dim (int) –

    Number of tensor modes (dimensions).

  • seed (int | None, default: None ) –

    Random seed for reproducible initialization.

  • dtype (dtype, default: jnp.float32 ) –

    Data type for all cores.

Raises:

  • ValueError –

    If d_dim, rank, or any entry of m_order is not a positive integer, or if the length of m_order does not match d_dim.

References

[1] G. Ballard, T.G. Kolda, "Tensor Decompositions for Data Science", Cambridge University Press, 2025. Part III: CP Decomposition.

[2] A. Cichocki, N. Lee, I. Oseledets, A-H. Phan, Q. Zhao, D.P. Mandic, "Tensor Networks for Dimensionality Reduction and Large-scale Optimization: Part I: Low-Rank Tensor Decompositions", Foundations and Trends In Machine Learning, 2016.

Examples:

>>> from tnkm.tn import CPNetwork
>>> tn = CPNetwork(m_order=2, rank=3, d_dim=4, seed=0)
>>> len(tn)
4
>>> tn.w_ten[0].shape
(2, 3)
>>> tn.w_vec.shape
(24,)
>>> tn.w_shapes
((2, 3), (2, 3), (2, 3), (2, 3))

w_vec property 

w_vec: Array

Vectorize all tensor network cores into a single parameter vector.

All cores in self.w_ten are flattened in Fortran (column-major) order and concatenated into a 1D vector. This ordering is important for consistency with tensor unfolding and optimization routines.

Returns:

  • Array –

    1D vector containing all tensor network parameters.

w_shapes property 

w_shapes: tuple[tuple[int, ...], ...]

Return the shapes of all tensor network cores.

Each entry corresponds to the shape of a core tensor in self.w_ten. The ordering matches the internal core ordering of the tensor network.

Returns:

  • tuple of tuple of int –

    Shapes of all cores, in network order.

set_cores 

set_cores(cores: list[Array]) -> None

Update CP decomposition factor matrices (cores).

Parameters:

  • cores (list[Array]) –

    List of factor matrices, each of shape (m_i, rank). The list must have length equal to d_dim, and the rank dimension must be consistent across all cores.

Raises:

  • ValueError –

    If the number of cores does not match d_dim, if any core has incorrect shape, or if the CP rank is inconsistent across cores.

Examples:

>>> import jax.numpy as jnp
>>> from tnkm.tn import CPNetwork
>>> tn = CPNetwork(m_order=2, rank=2, d_dim=2, seed=0)
>>> new_core_1 = jnp.arange(4).reshape((2, 2))
>>> new_core_2 = jnp.arange(4).reshape((2, 2)) + 10
>>> tn.set_cores([new_core_1, new_core_2])
>>> tn.w_ten
[Array([[0., 1.],
        [2., 3.]], dtype=float32),
 Array([[10., 11.],
        [12., 13.]], dtype=float32)]

TTNetwork dataclass 

TTNetwork(
    m_order: int | Sequence[int],
    rank: int | Sequence[int],
    d_dim: int,
    seed: int | None = None,
    dtype: dtype = jnp.float32,
)

Bases: TensorNetwork

Tensor Train (TT) tensor network representation.

Represents a high-order tensor using the Tensor Train (TT) decomposition as a chain of interconnected low-order cores. Each core has shape (r_i, m_i, r_{i+1}), where m_i is the local mode dimension and r_i are the TT ranks.

Parameters:

  • m_order (int | Sequence[int]) –

    Local dimension per mode. If a single integer is provided, it is broadcast to all d_dim modes.

  • rank (int | Sequence[int]) –

    TT ranks. If a single integer r is provided, the ranks are [1, r, ..., r, 1]. If a sequence is provided, it must have length d_dim + 1 with the first and last entries equal to 1.

  • d_dim (int) –

    Number of tensor modes (dimensions).

  • seed (int | None, default: None ) –

    Random seed for reproducible initialization.

  • dtype (dtype, default: jnp.float32 ) –

    Data type for all cores.

Raises:

  • ValueError –

    If d_dim is not positive, if m_order or rank contain non-positive values, if their lengths are inconsistent with d_dim, or if the boundary TT ranks are not equal to 1.

References

[1] I. Oseledets, "Tensor-Train Decomposition", SIAM Journal on Scientific Computing, 2011.

[2] A. Cichocki, N. Lee, I. Oseledets, A-H. Phan, Q. Zhao, D.P. Mandic, "Tensor Networks for Dimensionality Reduction and Large-scale Optimization: Part I: Low-Rank Tensor Decompositions", Foundations and Trends In Machine Learning, 2016.

Examples:

>>> from tnkm.tn import TTNetwork
>>> tn = TTNetwork(m_order=(2, 3, 4), rank=5, d_dim=3, seed=0)
>>> len(tn)
3
>>> tn.w_ten[0].shape
(1, 2, 5)
>>> tn.w_vec.shape
(105,)
>>> tn.w_shapes
((1, 2, 5), (5, 3, 5), (5, 4, 1))

w_vec property 

w_vec: Array

Vectorize all tensor network cores into a single parameter vector.

All cores in self.w_ten are flattened in Fortran (column-major) order and concatenated into a 1D vector. This ordering is important for consistency with tensor unfolding and optimization routines.

Returns:

  • Array –

    1D vector containing all tensor network parameters.

w_shapes property 

w_shapes: tuple[tuple[int, ...], ...]

Return the shapes of all tensor network cores.

Each entry corresponds to the shape of a core tensor in self.w_ten. The ordering matches the internal core ordering of the tensor network.

Returns:

  • tuple of tuple of int –

    Shapes of all cores, in network order.

set_cores 

set_cores(cores: list[Array]) -> None

Update the Tensor Train cores.

Parameters:

  • cores (list[Array]) –

    List of TT cores, where the i-th core has shape (r_i, m_i, r_{i+1}). The list must have length d_dim, and the TT ranks must be consistent between adjacent cores.

Raises:

  • ValueError –

    If the number of cores does not match d_dim, if any core has an invalid shape, or if adjacent TT ranks are inconsistent.

Examples:

>>> import jax.numpy as jnp
>>> from tnkm.tn import TTNetwork
>>> tn = TTNetwork(m_order=2, rank=2, d_dim=2, seed=0)
>>> new_core_1 = jnp.arange(4).reshape((1, 2, 2))
>>> new_core_2 = jnp.arange(4).reshape((2, 2, 1)) + 10
>>> tn.set_cores([new_core_1, new_core_2])
>>> tn.w_ten
[Array([[[0., 1.],
        [2., 3.]]], dtype=float32),
 Array([[[10.],
         [11.]],
        [[12.],
         [13.]]], dtype=float32)]