besskge.bess

PyTorch modules implementing the BESS distribution scheme [CJM+22] for KGE training and inference on multiple IPUs.

Classes

AllScoresBESS(candidate_sampler, score_fn[, ...])	Distributed scoring of (h, r, ?) or (?, r, t) queries against the entities in the knowledge graph, returning all scores to host in blocks, based on the BESS [CJM+22] inference scheme.
BessKGE(negative_sampler, score_fn[, ...])	Base class for distributed training and inference of KGE models, using the distribution framework BESS [CJM+22].
EmbeddingMovingBessKGE(negative_sampler, ...)	Compute negative scores on the shard where the positive triples are scored (namely the head shard).
ScoreMovingBessKGE(negative_sampler, score_fn)	Compute negative scores on the shard where the negative entities are stored.
TopKQueryBessKGE(k, candidate_sampler, score_fn)	Distributed scoring of (h, r, ?) or (?, r, t) queries (against all entities in the knowledge graph, or a query-specific set) returning the top-k most likely completions, based on the BESS [CJM+22] inference scheme.