TensorTrainRepresentation

class TensorTrainRepresentation(assignment: Tensor | None = None, num_cores: int = 3, ranks: int | Sequence[int] = 2, bases: str | X | type[X] | None | Sequence[str | X | type[X] | None] = None, bases_kwargs: Mapping[str, Any] | None | Sequence[Mapping[str, Any] | None] = 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:

  • assignment (LongTensor | None) – shape: (max_id, num_cores) the assignment on each level

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

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

  • bases (OneOrManyHintOrType) – the base representations for each level, or hints thereof.

  • bases_kwargs (OneOrManyOptionalKwargs) – 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: Tensor, max_id: int, num_cores: int, ms: Sequence[int])[source]

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: Sequence[int], ns: Sequence[int], max_id: int, shape: tuple[int, ...], num_cores: int)[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:

  • 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: int, num_cores: int, ms: Sequence[int]) → Tensor[source]

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

Returns:

shape: (max_id, num_cores) the assignment

Return type:

Tensor

classmethod factor_sizes(max_id: int, shape: Sequence[int], num_cores: int) → tuple[Sequence[int], Sequence[int]][source]

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

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} \]

Return type:

tuple[Sequence[int], Sequence[int]]

static get_shapes_and_einsum_eq(ranks: Sequence[int], ns: Sequence[int]) → tuple[str, Sequence[tuple[int, ...]]][source]

Determine core shapes and einsum equation.

Parameters:
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.

Return type:

tuple[str, Sequence[tuple[int, …]]]

iter_extra_repr() → Iterable[str][source]

Iterate over components for extra_repr().

Return type:

Iterable[str]