TensorTrainRepresentation

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

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:

• assignment (Optional[LongTensor]) – shape: (max_id, num_cores) the assignment on each level

• num_cores (int) – the number of cores to use

• ranks (Union[int, Sequence[int]]) – length: num_cores - 1 the individual ranks. Note that \(R_0 = R_d = 1\) should not be included

• bases (Union[str, ~X, Type[~X], None, Sequence[Union[str, ~X, Type[~X], None]]]) – the base representations for each level, or hints thereof.

• bases_kwargs (Union[Mapping[str, Any], None, Sequence[Optional[Mapping[str, Any]]]]) – keyword-based parameters for the bases

• kwargs – additional keyword-based parameters passed to Representation.__init__()

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.
iter_extra_repr()	Iterate over components for extra_repr().

Methods Documentation

static check_assignment(assignment, max_id, num_cores, ms)

Check that the assignment matches the other properties.

Parameters:

• assignment (Tensor) – shape: (max_id, num_cores) the assignment

• max_id (int) – the number of representations

• num_cores (int) – the number of tensor-train cores

• ms (Sequence[int]) – the individual sizes \(m_i\)

Raises:

ValueError – if the assignment is invalid

static check_factors(ms, ns, max_id, shape, num_cores)

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:

• ms (Sequence[int]) – length: num_cores the \(M\) factors \(m_i\)

• ns (Sequence[int]) – length: num_cores the \(N\) factors \(n_i\)

• max_id (int) – the maximum id, \(M\)

• shape (Tuple[int, …]) – the shape, \(N=prod(shape)\)

• num_cores (int) – the number of cores

Raises:

ValueError – if any of the conditions is violated

static create_default_assignment(max_id, num_cores, ms)

Create an assignment without using structural information.

Parameters:

• max_id (int) – the number of representations

• num_cores (int) – the number of tensor cores

• ms (Sequence[int]) – the sizes \(m_i\)

Return type:

LongTensor

Returns:

shape: (max_id, num_cores) the assignment

classmethod factor_sizes(max_id, shape, num_cores)

Factor the representation shape into smaller shapes for the cores.

Parameters:

• max_id (int) – the number of representations, “row count”, \(M\)

• shape (Sequence[int]) – the shape of an individual representation, “column count”, \(N\)

• num_cores (int) – the number of cores, \(k\)

Return type:

Tuple[Sequence[int], Sequence[int]]

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)

Determine core shapes and einsum equation.

Parameters:

• ranks (Sequence[int]) – the core ranks

• ns (Sequence[int]) – the sizes \(n_i\)

Return type:

Tuple[str, Sequence[Tuple[int, …]]]

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()

Iterate over components for extra_repr().

Return type:

Iterable[str]