nsl.lib.replicate_embeddings

Replicates the given embeddings by replicate_times.

nsl.lib.replicate_embeddings(
    embeddings, replicate_times
)

This function is useful when comparing the same instance with multiple other instances. For example, given a seed and its neighbors, this function can be used to replicate the embeddings of the seed by the number of its neighbors, such that the distances between the seed and its neighbors can be computed efficiently.

The replicate_times argument is either a scalar, or a 1-D tensor. For example, if

embeddings = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]

then we would have the following results for different replicate_times arguments:

replicate_times = 2
result = [[0, 1, 2], [0, 1, 2], [3, 4, 5], [3, 4, 5], [6, 7, 8], [6, 7, 8]]

and

replicate_times = [3, 0, 1]
result = [[0, 1, 2], [0, 1, 2], [0, 1, 2], [6, 7, 8]]

Args
embeddings	A Tensor of shape [batch_size, d1, ..., dN].
replicate_times	An integer scalar or an integer 1-D Tensor of shape [batch size]. Each element indicates the number of times the corresponding row in embeddings should be replicated.

Returns
A Tensor of shape [N, d1, ..., dN], where N is the sum of all elements in replicate_times.

Raises
InvalidArgumentError	If any value in replicate_times is negative.
TypeError	If replicate_times contains any value that cannot be cast to the int32 type.