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:
assignment (
Optional
[LongTensor
]) – shape: (max_id, num_cores) the assignment on each levelnum_cores (
int
) – the number of cores to useranks (
Union
[int
,Sequence
[int
]]) – length: num_cores - 1 the individual ranks. Note that \(R_0 = R_d = 1\) should not be includedbases (
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 baseskwargs – 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.
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.
- 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.
- 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.