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 train factorization of representations.

In the simple case without provided assignment this corresponds to TT-emb described in [yin2022].

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.

See also

Initialize the representation.

Parameters:

  • assignment (LongTensor | None) – shape: (max_id, num_cores) The core-assignment for each index on each level. If None, create_default_assignment() is used.

  • num_cores (int) – The number of cores to use.

  • ranks (OneOrSequence[int]) – length: num_cores - 1 The individual ranks for each core. 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

Raises:

ValueError – If the input validation on ranks or assignment failed.

Note

The parameter pair (bases, bases_kwargs) is used for pykeen.nn.representation_resolver

An explanation of resolvers and how to use them is given in https://class-resolver.readthedocs.io/en/latest/.

Methods Summary

check_assignment(assignment, max_id, ...)

Check that the assignment match in shape and its values are valid core "row" indices.

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]) → None[source]

Check that the assignment match in shape and its values are valid core “row” indices.

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.

Return type:

None

static check_factors(ms: Sequence[int], ns: Sequence[int], max_id: int, shape: tuple[int, ...], num_cores: int) → None[source]

Check whether the factors match the other parts.

Verifies that

\[M \leq \prod \limits_{m_i \in \textit{ms}} m_i \quad N \leq \prod \limits_{n_i \in \textit{ns}} n_i\]
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.

Return type:

None

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.

Note

This method implements a very simple heuristic of using the same value for each \(m_i\) / \(n_i\).

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

\[M \leq \prod \limits_{m_i \in \textit{ms}} m_i \quad N \leq \prod \limits_{n_i \in \textit{ns}} n_i\]

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]