# TensorTrainRepresentation¶

class TensorTrainRepresentation(assignment=None, num_cores=3, ranks=2, bases=None, bases_kwargs=None, **kwargs)[source]

Bases: Representation

A tensor factorization of representations.

In the simple case without provided assignment this corresponds to TT-emb described in https://assets.amazon.science/5c/0f/dd3eb08c4df88f2b4722e5fa8a7c/nimble-gnn-embedding-with-tensor-train-decomposition.pdf

where

$\mathbf{A}[i_1 \cdot \ldots \cdot i_k, j_1 \cdot \ldots \cdot j_k] = \sum_{r_i, \ldots, r_k} \mathbf{G}_1[0, i_1, j_1, r_1] \cdot \mathbf{G}_2[r_1, i_2, j_2, r_2] \cdot \ldots \cdot \mathbf{G}_k[r_k, i_k, j_k, 0]$

with TT core $$\mathbf{G}_i$$ of shape $$R_{i-1} \times m_i \times n_i \times R_i$$ and $$R_0 = R_d = 1$$.

Another variant in the paper used an assignment based on hierarchical topological clustering.

Initialize the representation.

Parameters:
Raises:

ValueError – if the input validation on ranks or assignment failed

Methods Summary

 check_assignment(assignment, max_id, ...) Check that the assignment matches the other properties. check_factors(ms, ns, max_id, shape, num_cores) Check whether the factors match the other parts. create_default_assignment(max_id, num_cores, ms) Create an assignment without using structural information. factor_sizes(max_id, shape, num_cores) Factor the representation shape into smaller shapes for the cores. get_shapes_and_einsum_eq(ranks, ns) Determine core shapes and einsum equation. Iterate over components for extra_repr().

Methods Documentation

static check_assignment(assignment, max_id, num_cores, ms)[source]

Check that the assignment matches the other properties.

Parameters:
Raises:

ValueError – if the assignment is invalid

static check_factors(ms, ns, max_id, shape, num_cores)[source]

Check whether the factors match the other parts.

Verifies that

$\prod \limits_{m_i \in ms} m_i \geq M \prod \limits_{n_i \in ns} n_i \geq N$
Parameters:
Raises:

ValueError – if any of the conditions is violated

static create_default_assignment(max_id, num_cores, ms)[source]

Create an assignment without using structural information.

Parameters:
Return type:

LongTensor

Returns:

shape: (max_id, num_cores) the assignment

classmethod factor_sizes(max_id, shape, num_cores)[source]

Factor the representation shape into smaller shapes for the cores.

Parameters:
Return type:
Returns:

a tuple (ms, ns) of positive integer sequences of length $$k$$ fulfilling

\begin{align}\begin{aligned}\prod \limits_{m_i \in ms} m_i \geq M\\\prod \limits_{n_i \in ns} n_i \geq N\end{aligned}\end{align}

static get_shapes_and_einsum_eq(ranks, ns)[source]

Determine core shapes and einsum equation.

Parameters:
Return type:
Returns:

a pair (eq, shapes), where eq is a valid einsum equation and shapes a sequence of representation shapes. Notice that the shapes do not include the “max_id dimension” of the resulting embedding.

iter_extra_repr()[source]

Iterate over components for extra_repr().

Return type: