besskge.embedding.initialize_entity_embedding

besskge.embedding.initialize_entity_embedding(sharding, initializer, row_size=None)

Initialize entity embedding table.

Parameters:

• sharding (Sharding) – Entity sharding.

• initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Embedding table or list of initializing functions. If providing an embedding table, this can either be sharded (shape: [n_shard, max_entity_per_shard, row_size]) or unsharded [shape: (n_entity, row_size]). If providing list of initializers, this needs to be of same length as row_size.

• row_size (Optional[List[int]]) – Number of parameters for each entity. This needs to be a list, with the lengths of the different embedding tensors to allocate for each entity. Each embedding tensor, once allocated, is initialized with the corresponding entry of initializer. Can be omitted if passing an embedding table as initializer.

Return type:

Parameter

Returns:

shape: (n_shard, max_ent_per_shard, row_size) Entity embedding table.