diff --git a/_modules/besskge/utils.html b/_modules/besskge/utils.html index a250d2d..28b723d 100644 --- a/_modules/besskge/utils.html +++ b/_modules/besskge/utils.html @@ -140,7 +140,7 @@

Source code for besskge.utils

         Row-wise rotated tensors.
     """
     # Always compute sin and cos in fp16, as faster on IPU
-    if r.dtype == torch.float32:
+    if r.dtype == torch.float32 and r.device.type == "ipu":
         r_cos = torch.cos(r.to(dtype=torch.float16)).to(dtype=torch.float32)
         r_sin = torch.sin(r.to(dtype=torch.float16)).to(dtype=torch.float32)
     else:
diff --git a/generated/besskge.batch_sampler.RandomShardedBatchSampler.html b/generated/besskge.batch_sampler.RandomShardedBatchSampler.html
index 8473af3..fd4c3c3 100644
--- a/generated/besskge.batch_sampler.RandomShardedBatchSampler.html
+++ b/generated/besskge.batch_sampler.RandomShardedBatchSampler.html
@@ -109,18 +109,18 @@ 
besskge.batch_sampler.RandomShardedBatchSampler
    
 
  • partitioned_triple_set (PartitionedTripleSet) – The pre-processed collection of triples.
    • 
 
  • negative_sampler (ShardedNegativeSampler) – The sampler for negative entities.
    • 
-
  • shard_bs (int) – The micro-batch size. This is the number of positive triples
+
  • shard_bs (int) – The micro-batch size. This is the number of positive triples
 processed on each shard.
    • 
-
  • batches_per_step (int) – The number of batches to sample at each call.
    • 
-
  • seed (int) – The RNG seed.
    • 
-
  • hrt_freq_weighting (bool) – If True, uses frequency-based triple weighting. Default: False.
    • 
-
  • weight_smoothing (float) – Weight-smoothing parameter for frequency-based
+
  • batches_per_step (int) – The number of batches to sample at each call.
    • 
+
  • seed (int) – The RNG seed.
    • 
+
  • hrt_freq_weighting (bool) – If True, uses frequency-based triple weighting. Default: False.
    • 
+
  • weight_smoothing (float) – Weight-smoothing parameter for frequency-based
 triple weighting. Default: 0.0.
    • 
-
  • duplicate_batch (bool) – If True, the batch sampled from each triple partition has two
+
  • duplicate_batch (bool) – If True, the batch sampled from each triple partition has two
 identical halves.
 This is to be used with “ht” corruption scheme at inference time.
 Default: False.
    • 
-
  • return_triple_idx (bool) – If True, return the indices (wrt partitioned_triple_set.triples)
+
  • return_triple_idx (bool) – If True, return the indices (wrt partitioned_triple_set.triples)
 of the triples in the batch. Default: False.
    • 
 

 
@@ -136,10 +136,10 @@ 
besskge.batch_sampler.RandomShardedBatchSamplerParameters:
 
    
 
  • options (Options) – poptorch.Options used to compile and run the model.
    • 
-
  • shuffle (bool) – If True, shuffles triples at each new epoch. Default: True.
    • 
-
  • num_workers (int) – see torch.utils.data.DataLoader.__init__(). Default: 0.
    • 
-
  • persistent_workers (bool) – see torch.utils.data.DataLoader.__init__(). Default: False.
    • 
-
  • buffer_size (int) – Size of the ring buffer in shared memory used to preload batches.
    • 
+
  • shuffle (bool) – If True, shuffles triples at each new epoch. Default: True.
    • 
+
  • num_workers (int) – see torch.utils.data.DataLoader.__init__(). Default: 0.
    • 
+
  • persistent_workers (bool) – see torch.utils.data.DataLoader.__init__(). Default: False.
    • 
+
  • buffer_size (int) – Size of the ring buffer in shared memory used to preload batches.
    • 
 

 

 
Return type:

@@ -156,14 +156,14 @@ 
besskge.batch_sampler.RandomShardedBatchSamplerget_dataloader_sampler(shuffle=True)[source]
 
Returns the dataloader sampler.

 
Instantiate the appropriate torch.data.Sampler class for the
-torch.utils.data.DataLoader class to be used with the
+torch.utils.data.DataLoader class to be used with the
 sharded batch sampler.

 

 
Parameters:

-
shuffle (bool) – Shuffle triples at each new epoch.

+
shuffle (bool) – Shuffle triples at each new epoch.

 

 
Return type:

-
Sampler[List[int]]

+
Sampler[List[int]]

 

 
Returns:

 
The dataloader sampler.

@@ -177,10 +177,10 @@ 
besskge.batch_sampler.RandomShardedBatchSampler
 
Parameters:

-
idx (List[int]) – The batch index.

+
idx (List[int]) – The batch index.

 

 
Return type:

-
Dict[str, Union[ndarray[Any, dtype[int64]], ndarray[Any, dtype[bool_]]]]

+
Dict[str, Union[ndarray[Any, dtype[int64]], ndarray[Any, dtype[bool_]]]]

 

 
Returns:

 
Per-partition indices of positive triples, and other relevant data.

@@ -192,13 +192,13 @@ 
besskge.batch_sampler.RandomShardedBatchSampler
 static worker_init_fn(worker_id)
 
Worker initialization function to be passed to
-torch.utils.data.DataLoader.

+torch.utils.data.DataLoader.

 

 
Parameters:

-
worker_id (int) – Worker ID.

+
worker_id (int) – Worker ID.

 

 
Return type:

-
None

+
None

 

 

 

diff --git a/generated/besskge.batch_sampler.RigidShardedBatchSampler.html b/generated/besskge.batch_sampler.RigidShardedBatchSampler.html
index 2d1cf52..bb04c28 100644
--- a/generated/besskge.batch_sampler.RigidShardedBatchSampler.html
+++ b/generated/besskge.batch_sampler.RigidShardedBatchSampler.html
@@ -111,18 +111,18 @@ 
besskge.batch_sampler.RigidShardedBatchSampler
    
 
  • partitioned_triple_set (PartitionedTripleSet) – The pre-processed collection of triples.
    • 
 
  • negative_sampler (ShardedNegativeSampler) – The sampler for negative entities.
    • 
-
  • shard_bs (int) – The micro-batch size. This is the number of positive triples
+
  • shard_bs (int) – The micro-batch size. This is the number of positive triples
 processed on each shard.
    • 
-
  • batches_per_step (int) – The number of batches to sample at each call.
    • 
-
  • seed (int) – The RNG seed.
    • 
-
  • hrt_freq_weighting (bool) – If True, uses frequency-based triple weighting. Default: False.
    • 
-
  • weight_smoothing (float) – Weight-smoothing parameter for frequency-based
+
  • batches_per_step (int) – The number of batches to sample at each call.
    • 
+
  • seed (int) – The RNG seed.
    • 
+
  • hrt_freq_weighting (bool) – If True, uses frequency-based triple weighting. Default: False.
    • 
+
  • weight_smoothing (float) – Weight-smoothing parameter for frequency-based
 triple weighting. Default: 0.0.
    • 
-
  • duplicate_batch (bool) – If True, the batch sampled from each triple partition has two
+
  • duplicate_batch (bool) – If True, the batch sampled from each triple partition has two
 identical halves.
 This is to be used with “ht” corruption scheme at inference time.
 Default: False.
    • 
-
  • return_triple_idx (bool) – If True, return the indices (wrt partitioned_triple_set.triples)
+
  • return_triple_idx (bool) – If True, return the indices (wrt partitioned_triple_set.triples)
 of the triples in the batch. Default: False.
    • 
 

 

@@ -138,10 +138,10 @@ 
besskge.batch_sampler.RigidShardedBatchSamplerParameters:
 
    
 
  • options (Options) – poptorch.Options used to compile and run the model.
    • 
-
  • shuffle (bool) – If True, shuffles triples at each new epoch. Default: True.
    • 
-
  • num_workers (int) – see torch.utils.data.DataLoader.__init__(). Default: 0.
    • 
-
  • persistent_workers (bool) – see torch.utils.data.DataLoader.__init__(). Default: False.
    • 
-
  • buffer_size (int) – Size of the ring buffer in shared memory used to preload batches.
    • 
+
  • shuffle (bool) – If True, shuffles triples at each new epoch. Default: True.
    • 
+
  • num_workers (int) – see torch.utils.data.DataLoader.__init__(). Default: 0.
    • 
+
  • persistent_workers (bool) – see torch.utils.data.DataLoader.__init__(). Default: False.
    • 
+
  • buffer_size (int) – Size of the ring buffer in shared memory used to preload batches.
    • 
 

 

 
Return type:

@@ -158,14 +158,14 @@ 
besskge.batch_sampler.RigidShardedBatchSamplerget_dataloader_sampler(shuffle)
 
Returns the dataloader sampler.

 
Instantiate the appropriate torch.data.Sampler class for the
-torch.utils.data.DataLoader class to be used with the
+torch.utils.data.DataLoader class to be used with the
 sharded batch sampler.

 

 
Parameters:

-
shuffle (bool) – Shuffle triples at each new epoch.

+
shuffle (bool) – Shuffle triples at each new epoch.

 

 
Return type:

-
Sampler[List[int]]

+
Sampler[List[int]]

 

 
Returns:

 
The dataloader sampler.

@@ -179,10 +179,10 @@ 
besskge.batch_sampler.RigidShardedBatchSampler
 
Parameters:

-
idx (List[int]) – The batch index.

+
idx (List[int]) – The batch index.

 

 
Return type:

-
Dict[str, Union[ndarray[Any, dtype[int64]], ndarray[Any, dtype[bool_]]]]

+
Dict[str, Union[ndarray[Any, dtype[int64]], ndarray[Any, dtype[bool_]]]]

 

 
Returns:

 
Per-partition indices of positive triples, and other relevant data.

@@ -194,13 +194,13 @@ 
besskge.batch_sampler.RigidShardedBatchSampler
 static worker_init_fn(worker_id)
 
Worker initialization function to be passed to
-torch.utils.data.DataLoader.

+torch.utils.data.DataLoader.

 

 
Parameters:

-
worker_id (int) – Worker ID.

+
worker_id (int) – Worker ID.

 

 
Return type:

-
None

+
None

 

 

 

diff --git a/generated/besskge.batch_sampler.ShardedBatchSampler.html b/generated/besskge.batch_sampler.ShardedBatchSampler.html
index 2289457..130c293 100644
--- a/generated/besskge.batch_sampler.ShardedBatchSampler.html
+++ b/generated/besskge.batch_sampler.ShardedBatchSampler.html
@@ -108,18 +108,18 @@ 
besskge.batch_sampler.ShardedBatchSampler
    
 
  • partitioned_triple_set (PartitionedTripleSet) – The pre-processed collection of triples.
    • 
 
  • negative_sampler (ShardedNegativeSampler) – The sampler for negative entities.
    • 
-
  • shard_bs (int) – The micro-batch size. This is the number of positive triples
+
  • shard_bs (int) – The micro-batch size. This is the number of positive triples
 processed on each shard.
    • 
-
  • batches_per_step (int) – The number of batches to sample at each call.
    • 
-
  • seed (int) – The RNG seed.
    • 
-
  • hrt_freq_weighting (bool) – If True, uses frequency-based triple weighting. Default: False.
    • 
-
  • weight_smoothing (float) – Weight-smoothing parameter for frequency-based
+
  • batches_per_step (int) – The number of batches to sample at each call.
    • 
+
  • seed (int) – The RNG seed.
    • 
+
  • hrt_freq_weighting (bool) – If True, uses frequency-based triple weighting. Default: False.
    • 
+
  • weight_smoothing (float) – Weight-smoothing parameter for frequency-based
 triple weighting. Default: 0.0.
    • 
-
  • duplicate_batch (bool) – If True, the batch sampled from each triple partition has two
+
  • duplicate_batch (bool) – If True, the batch sampled from each triple partition has two
 identical halves.
 This is to be used with “ht” corruption scheme at inference time.
 Default: False.
    • 
-
  • return_triple_idx (bool) – If True, return the indices (wrt partitioned_triple_set.triples)
+
  • return_triple_idx (bool) – If True, return the indices (wrt partitioned_triple_set.triples)
 of the triples in the batch. Default: False.
    • 
 

 

@@ -135,10 +135,10 @@ 
besskge.batch_sampler.ShardedBatchSamplerParameters:
 
    
 
  • options (Options) – poptorch.Options used to compile and run the model.
    • 
-
  • shuffle (bool) – If True, shuffles triples at each new epoch. Default: True.
    • 
-
  • num_workers (int) – see torch.utils.data.DataLoader.__init__(). Default: 0.
    • 
-
  • persistent_workers (bool) – see torch.utils.data.DataLoader.__init__(). Default: False.
    • 
-
  • buffer_size (int) – Size of the ring buffer in shared memory used to preload batches.
    • 
+
  • shuffle (bool) – If True, shuffles triples at each new epoch. Default: True.
    • 
+
  • num_workers (int) – see torch.utils.data.DataLoader.__init__(). Default: 0.
    • 
+
  • persistent_workers (bool) – see torch.utils.data.DataLoader.__init__(). Default: False.
    • 
+
  • buffer_size (int) – Size of the ring buffer in shared memory used to preload batches.
    • 
 

 

 
Return type:

@@ -155,14 +155,14 @@ 
besskge.batch_sampler.ShardedBatchSamplerget_dataloader_sampler(shuffle)[source]
 
Returns the dataloader sampler.

 
Instantiate the appropriate torch.data.Sampler class for the
-torch.utils.data.DataLoader class to be used with the
+torch.utils.data.DataLoader class to be used with the
 sharded batch sampler.

 

 
Parameters:

-
shuffle (bool) – Shuffle triples at each new epoch.

+
shuffle (bool) – Shuffle triples at each new epoch.

 

 
Return type:

-
Sampler[List[int]]

+
Sampler[List[int]]

 

 
Returns:

 
The dataloader sampler.

@@ -176,10 +176,10 @@ 
besskge.batch_sampler.ShardedBatchSampler
 
Parameters:

-
idx (List[int]) – The batch index.

+
idx (List[int]) – The batch index.

 

 
Return type:

-
Dict[str, Union[ndarray[Any, dtype[int64]], ndarray[Any, dtype[bool_]]]]

+
Dict[str, Union[ndarray[Any, dtype[int64]], ndarray[Any, dtype[bool_]]]]

 

 
Returns:

 
Per-partition indices of positive triples, and other relevant data.

@@ -191,13 +191,13 @@ 
besskge.batch_sampler.ShardedBatchSampler
 static worker_init_fn(worker_id)[source]
 
Worker initialization function to be passed to
-torch.utils.data.DataLoader.

+torch.utils.data.DataLoader.

 

 
Parameters:

-
worker_id (int) – Worker ID.

+
worker_id (int) – Worker ID.

 

 
Return type:

-
None

+
None

 

 

 

diff --git a/generated/besskge.bess.BessKGE.html b/generated/besskge.bess.BessKGE.html
index 8eea5aa..b745a7e 100644
--- a/generated/besskge.bess.BessKGE.html
+++ b/generated/besskge.bess.BessKGE.html
@@ -112,12 +112,12 @@ 
besskge.bess.BessKGE
    
 
  • negative_sampler (ShardedNegativeSampler) – Sampler of negative entities.
    • 
 
  • score_fn (BaseScoreFunction) – Scoring function.
    • 
-
  • loss_fn (Optional[BaseLossFunction]) – Loss function, required when training. Default: None.
    • 
-
  • evaluation (Optional[Evaluation]) – Evaluation module, for computing metrics on device.
+
  • loss_fn (Optional[BaseLossFunction]) – Loss function, required when training. Default: None.
    • 
+
  • evaluation (Optional[Evaluation]) – Evaluation module, for computing metrics on device.
 Default: None.
    • 
-
  • return_scores (bool) – If True, return positive and negative scores of batches to the host.
+
  • return_scores (bool) – If True, return positive and negative scores of batches to the host.
 Default: False.
    • 
-
  • augment_negative (bool) – If True, augment sampled negative entities with the head/tails
+
  • augment_negative (bool) – If True, augment sampled negative entities with the head/tails
 (according to the corruption scheme) of other positive triples
 in the micro-batch. Default: False.
    • 
 

@@ -138,26 +138,26 @@ 
besskge.bess.BessKGE
 
Parameters:

 
    
-
  • head (Tensor) – shape: (1, n_shard, positive_per_partition)
+
  • head (Tensor) – shape: (1, n_shard, positive_per_partition)
 Head indices.
    • 
-
  • relation (Tensor) – shape: (1, n_shard, positive_per_partition)
+
  • relation (Tensor) – shape: (1, n_shard, positive_per_partition)
 Relation indices.
    • 
-
  • tail (Tensor) – shape: (1, n_shard, positive_per_partition)
+
  • tail (Tensor) – shape: (1, n_shard, positive_per_partition)
 Tail indices.
    • 
-
  • triple_mask (Optional[Tensor]) – shape: (1, n_shard, positive_per_partition)
+
  • triple_mask (Optional[Tensor]) – shape: (1, n_shard, positive_per_partition)
 Mask to filter the triples in the micro-batch
 before computing metrics.
    • 
-
  • negative (Tensor) – shape: (1, n_shard, B, padded_negative)
+
  • negative (Tensor) – shape: (1, n_shard, B, padded_negative)
 Indices of negative entities,
 with B = 1, 2 or n_shard * positive_per_partition.
    • 
-
  • triple_weight (Optional[Tensor]) – shape: (1, n_shard * positive_per_partition,) or (1,)
+
  • triple_weight (Optional[Tensor]) – shape: (1, n_shard * positive_per_partition,) or (1,)
 Weights of positive triples.
    • 
-
  • negative_mask (Optional[Tensor]) – shape: (1, B, n_shard, padded_negative)
+
  • negative_mask (Optional[Tensor]) – shape: (1, B, n_shard, padded_negative)
 Mask to identify padding negatives, to discard when computing metrics.
    • 
 

 

 
Return type:

-
Dict[str, Any]

+
Dict[str, Any]

 

 
Returns:

 
Micro-batch loss, scores and metrics.

@@ -167,7 +167,7 @@ 
besskge.bess.BessKGE
 

-property n_embedding_parameters: int

+property n_embedding_parameters: int
 
Returns the number of trainable parameters in the embedding tables

 

 
@@ -178,14 +178,14 @@ 
besskge.bess.BessKGE
 
Parameters:

 

 

 
Return type:

-
Tuple[Tensor, Tensor]

+
Tuple[Tensor, Tensor]

 

 
Returns:

 
Positive (shape: (n_shard * positive_per_partition,))
diff --git a/generated/besskge.bess.EmbeddingMovingBessKGE.html b/generated/besskge.bess.EmbeddingMovingBessKGE.html
index 1e4e385..4d6d256 100644
--- a/generated/besskge.bess.EmbeddingMovingBessKGE.html
+++ b/generated/besskge.bess.EmbeddingMovingBessKGE.html
@@ -116,12 +116,12 @@ 
besskge.bess.EmbeddingMovingBessKGE
    
 
  • negative_sampler (ShardedNegativeSampler) – Sampler of negative entities.
    • 
 
  • score_fn (BaseScoreFunction) – Scoring function.
    • 
-
  • loss_fn (Optional[BaseLossFunction]) – Loss function, required when training. Default: None.
    • 
-
  • evaluation (Optional[Evaluation]) – Evaluation module, for computing metrics on device.
+
  • loss_fn (Optional[BaseLossFunction]) – Loss function, required when training. Default: None.
    • 
+
  • evaluation (Optional[Evaluation]) – Evaluation module, for computing metrics on device.
 Default: None.
    • 
-
  • return_scores (bool) – If True, return positive and negative scores of batches to the host.
+
  • return_scores (bool) – If True, return positive and negative scores of batches to the host.
 Default: False.
    • 
-
  • augment_negative (bool) – If True, augment sampled negative entities with the head/tails
+
  • augment_negative (bool) – If True, augment sampled negative entities with the head/tails
 (according to the corruption scheme) of other positive triples
 in the micro-batch. Default: False.
    • 
 

@@ -142,26 +142,26 @@ 
besskge.bess.EmbeddingMovingBessKGE
 
Parameters:

 
    
-
  • head (Tensor) – shape: (1, n_shard, positive_per_partition)
+
  • head (Tensor) – shape: (1, n_shard, positive_per_partition)
 Head indices.
    • 
-
  • relation (Tensor) – shape: (1, n_shard, positive_per_partition)
+
  • relation (Tensor) – shape: (1, n_shard, positive_per_partition)
 Relation indices.
    • 
-
  • tail (Tensor) – shape: (1, n_shard, positive_per_partition)
+
  • tail (Tensor) – shape: (1, n_shard, positive_per_partition)
 Tail indices.
    • 
-
  • triple_mask (Optional[Tensor]) – shape: (1, n_shard, positive_per_partition)
+
  • triple_mask (Optional[Tensor]) – shape: (1, n_shard, positive_per_partition)
 Mask to filter the triples in the micro-batch
 before computing metrics.
    • 
-
  • negative (Tensor) – shape: (1, n_shard, B, padded_negative)
+
  • negative (Tensor) – shape: (1, n_shard, B, padded_negative)
 Indices of negative entities,
 with B = 1, 2 or n_shard * positive_per_partition.
    • 
-
  • triple_weight (Optional[Tensor]) – shape: (1, n_shard * positive_per_partition,) or (1,)
+
  • triple_weight (Optional[Tensor]) – shape: (1, n_shard * positive_per_partition,) or (1,)
 Weights of positive triples.
    • 
-
  • negative_mask (Optional[Tensor]) – shape: (1, B, n_shard, padded_negative)
+
  • negative_mask (Optional[Tensor]) – shape: (1, B, n_shard, padded_negative)
 Mask to identify padding negatives, to discard when computing metrics.
    • 
 

 

 
Return type:

-
Dict[str, Any]

+
Dict[str, Any]

 

 
Returns:

 
Micro-batch loss, scores and metrics.

@@ -171,7 +171,7 @@ 
besskge.bess.EmbeddingMovingBessKGE
 

-property n_embedding_parameters: int

+property n_embedding_parameters: int
 
Returns the number of trainable parameters in the embedding tables

 

 
@@ -182,14 +182,14 @@ 
besskge.bess.EmbeddingMovingBessKGE
 
Parameters:

 

 

 
Return type:

-
Tuple[Tensor, Tensor]

+
Tuple[Tensor, Tensor]

 

 
Returns:

 
Positive (shape: (n_shard * positive_per_partition,))
diff --git a/generated/besskge.bess.ScoreMovingBessKGE.html b/generated/besskge.bess.ScoreMovingBessKGE.html
index 90e861b..00cc533 100644
--- a/generated/besskge.bess.ScoreMovingBessKGE.html
+++ b/generated/besskge.bess.ScoreMovingBessKGE.html
@@ -119,12 +119,12 @@ 
besskge.bess.ScoreMovingBessKGE
    
 
  • negative_sampler (ShardedNegativeSampler) – Sampler of negative entities.
    • 
 
  • score_fn (BaseScoreFunction) – Scoring function.
    • 
-
  • loss_fn (Optional[BaseLossFunction]) – Loss function, required when training. Default: None.
    • 
-
  • evaluation (Optional[Evaluation]) – Evaluation module, for computing metrics on device.
+
  • loss_fn (Optional[BaseLossFunction]) – Loss function, required when training. Default: None.
    • 
+
  • evaluation (Optional[Evaluation]) – Evaluation module, for computing metrics on device.
 Default: None.
    • 
-
  • return_scores (bool) – If True, return positive and negative scores of batches to the host.
+
  • return_scores (bool) – If True, return positive and negative scores of batches to the host.
 Default: False.
    • 
-
  • augment_negative (bool) – If True, augment sampled negative entities with the head/tails
+
  • augment_negative (bool) – If True, augment sampled negative entities with the head/tails
 (according to the corruption scheme) of other positive triples
 in the micro-batch. Default: False.
    • 
 

@@ -145,26 +145,26 @@ 
besskge.bess.ScoreMovingBessKGE
 
Parameters:

 
    
-
  • head (Tensor) – shape: (1, n_shard, positive_per_partition)
+
  • head (Tensor) – shape: (1, n_shard, positive_per_partition)
 Head indices.
    • 
-
  • relation (Tensor) – shape: (1, n_shard, positive_per_partition)
+
  • relation (Tensor) – shape: (1, n_shard, positive_per_partition)
 Relation indices.
    • 
-
  • tail (Tensor) – shape: (1, n_shard, positive_per_partition)
+
  • tail (Tensor) – shape: (1, n_shard, positive_per_partition)
 Tail indices.
    • 
-
  • triple_mask (Optional[Tensor]) – shape: (1, n_shard, positive_per_partition)
+
  • triple_mask (Optional[Tensor]) – shape: (1, n_shard, positive_per_partition)
 Mask to filter the triples in the micro-batch
 before computing metrics.
    • 
-
  • negative (Tensor) – shape: (1, n_shard, B, padded_negative)
+
  • negative (Tensor) – shape: (1, n_shard, B, padded_negative)
 Indices of negative entities,
 with B = 1, 2 or n_shard * positive_per_partition.
    • 
-
  • triple_weight (Optional[Tensor]) – shape: (1, n_shard * positive_per_partition,) or (1,)
+
  • triple_weight (Optional[Tensor]) – shape: (1, n_shard * positive_per_partition,) or (1,)
 Weights of positive triples.
    • 
-
  • negative_mask (Optional[Tensor]) – shape: (1, B, n_shard, padded_negative)
+
  • negative_mask (Optional[Tensor]) – shape: (1, B, n_shard, padded_negative)
 Mask to identify padding negatives, to discard when computing metrics.
    • 
 

 

 
Return type:

-
Dict[str, Any]

+
Dict[str, Any]

 

 
Returns:

 
Micro-batch loss, scores and metrics.

@@ -174,7 +174,7 @@ 
besskge.bess.ScoreMovingBessKGE
 

-property n_embedding_parameters: int

+property n_embedding_parameters: int
 
Returns the number of trainable parameters in the embedding tables

 

 
@@ -185,14 +185,14 @@ 
besskge.bess.ScoreMovingBessKGE
 
Parameters:

 

 

 
Return type:

-
Tuple[Tensor, Tensor]

+
Tuple[Tensor, Tensor]

 

 
Returns:

 
Positive (shape: (n_shard * positive_per_partition,))
diff --git a/generated/besskge.bess.TopKQueryBessKGE.html b/generated/besskge.bess.TopKQueryBessKGE.html
index 5cd771c..9e78de2 100644
--- a/generated/besskge.bess.TopKQueryBessKGE.html
+++ b/generated/besskge.bess.TopKQueryBessKGE.html
@@ -119,17 +119,17 @@ 
besskge.bess.TopKQueryBessKGE
 
Parameters:

 
    
-
  • k (int) – For each query return the top-k most likely predictions.
    • 
-
  • candidate_sampler (Union[TripleBasedShardedNegativeSampler, PlaceholderNegativeSampler]) – Sampler of candidate entities to score against queries.
+
  • k (int) – For each query return the top-k most likely predictions.
    • 
+
  • candidate_sampler (Union[TripleBasedShardedNegativeSampler, PlaceholderNegativeSampler]) – Sampler of candidate entities to score against queries.
 Use besskge.negative_sampler.PlaceholderNegativeSampler
 to score queries against all entities in the knowledge graph, avoiding
 unnecessary loading of negative entities on device.
    • 
 
  • score_fn (BaseScoreFunction) – Scoring function.
    • 
-
  • evaluation (Optional[Evaluation]) – Evaluation module, for computing metrics on device.
+
  • evaluation (Optional[Evaluation]) – Evaluation module, for computing metrics on device.
 Default: None.
    • 
-
  • return_scores (bool) – If True, return scores of the top-k best completions.
+
  • return_scores (bool) – If True, return scores of the top-k best completions.
 Default: False.
    • 
-
  • window_size (int) – Size of the sliding window, namely the number of negative entities
+
  • window_size (int) – Size of the sliding window, namely the number of negative entities
 scored against each query at each step of the on-device for-loop.
 Should be decreased with large batch sizes, to avoid an OOM error.
 Default: 100.
    • 
@@ -149,22 +149,22 @@ 
    besskge.bess.TopKQueryBessKGE
 
    Parameters:
    
 
      
-
    • relation (Tensor) – shape: (shard_bs,)
+
    • relation (Tensor) – shape: (shard_bs,)
 Relation indices.
      • 
-
    • head (Optional[Tensor]) – shape: (shard_bs,)
+
    • head (Optional[Tensor]) – shape: (shard_bs,)
 Head indices, if known. Default: None.
      • 
-
    • tail (Optional[Tensor]) – shape: (shard_bs,)
+
    • tail (Optional[Tensor]) – shape: (shard_bs,)
 Tail indices, if known. Default: None.
      • 
-
    • negative (Optional[Tensor]) – shape: (n_shard, B, padded_negative)
+
    • negative (Optional[Tensor]) – shape: (n_shard, B, padded_negative)
 Candidates to score against the queries.
 It can be the same set for all queries (B=1),
 or specific for each query in the batch (B=shard_bs).
 If None, score each query against all entities in the knowledge
 graph. Default: None.
      • 
-
    • triple_mask (Optional[Tensor]) – shape: (shard_bs,)
+
    • triple_mask (Optional[Tensor]) – shape: (shard_bs,)
 Mask to filter the triples in the micro-batch
 before computing metrics. Default: None.
      • 
-
    • negative_mask (Optional[Tensor]) – shape: (n_shard, B, padded_negative)
+
    • negative_mask (Optional[Tensor]) – shape: (n_shard, B, padded_negative)
 If candidates are provided, mask to discard padding
 negatives when computing best completions.
 Requires the use of mask_on_gather=True in the candidate
@@ -174,7 +174,7 @@ 
      besskge.bess.TopKQueryBessKGEReturn type:
-
      Dict[str, Any]
      
+
      Dict[str, Any]
      
 
      
 
 
    
diff --git a/generated/besskge.dataset.KGDataset.html b/generated/besskge.dataset.KGDataset.html
index 5704d97..9d7f198 100644
--- a/generated/besskge.dataset.KGDataset.html
+++ b/generated/besskge.dataset.KGDataset.html
@@ -103,14 +103,14 @@ 
    besskge.dataset.KGDataset
 
    Parameters:
    
 
    
 
    
 
@@ -124,7 +124,7 @@ 
    besskge.dataset.KGDataset
 
    Parameters:
    
-
    root (Path) – Local path to the dataset. If the dataset is not present in this
+
    root (Path) – Local path to the dataset. If the dataset is not present in this
 location, then it is downloaded and stored here.
    
 
    
 
    Return type:
    
@@ -146,7 +146,7 @@ 
    besskge.dataset.KGDataset
 
    Parameters:
    
-
    root (Path) – Local path to the dataset. If the dataset is not present in this
+
    root (Path) – Local path to the dataset. If the dataset is not present in this
 location, then it is downloaded and stored here.
    
 
    
 
    Return type:
    
@@ -169,7 +169,7 @@ 
    besskge.dataset.KGDataset
 
    Parameters:
    
-
    root (Path) – Local path to the dataset. If the dataset is not present in this
+
    root (Path) – Local path to the dataset. If the dataset is not present in this
 location, then it is downloaded and stored here.
    
 
    
 
    Return type:
    
@@ -194,7 +194,7 @@ 
    besskge.dataset.KGDataset
 
    Parameters:
    
-
    root (Path) – Local path to the dataset. If the dataset is not present in this
+
    root (Path) – Local path to the dataset. If the dataset is not present in this
 location, then it is downloaded and stored here.
    
 
    
 
    Return type:
    
@@ -208,7 +208,7 @@ 
    besskge.dataset.KGDataset
 
    
-entity_dict: Optional[List[str]] = None
    
+entity_dict: Optional[List[str]] = None
 
    Entity labels by ID; str[n_entity]
    
 
    
 
@@ -221,17 +221,17 @@ 
    besskge.dataset.KGDataset
 
    Parameters:
    
 
      
-
    • df (Union[DataFrame, Dict[str, DataFrame]]) – Pandas DataFrame of all triples in the knowledge graph dataset,
+
    • df (Union[DataFrame, Dict[str, DataFrame]]) – Pandas DataFrame of all triples in the knowledge graph dataset,
 or dictionary of DataFrames of triples for each part of the dataset split
      • 
-
    • head_column (Union[int, str]) – Name of the DataFrame column storing head entities
      • 
-
    • relation_column (Union[int, str]) – Name of the DataFrame column storing relations
      • 
-
    • tail_column (Union[int, str]) – Name of the DataFrame column storing tail entities
      • 
-
    • entity_types (Union[Series, Dict[str, str], None]) – If entities have types, dictionary or pandas Series of mappings
+
    • head_column (Union[int, str]) – Name of the DataFrame column storing head entities
      • 
+
    • relation_column (Union[int, str]) – Name of the DataFrame column storing relations
      • 
+
    • tail_column (Union[int, str]) – Name of the DataFrame column storing tail entities
      • 
+
    • entity_types (Union[Series, Dict[str, str], None]) – If entities have types, dictionary or pandas Series of mappings
 entity label -> entity type (as strings).
      • 
-
    • split (Tuple[float, float, float]) – Tuple to set the train/validation/test split.
+
    • split (Tuple[float, float, float]) – Tuple to set the train/validation/test split.
 Only used if no pre-defined dataset split is specified,
 i.e. if df is not a dictionary.
      • 
-
    • seed (int) – Random seed for the train/validation/test split.
+
    • seed (int) – Random seed for the train/validation/test split.
 Only used if no pre-defined dataset split is specified,
 i.e. if df is not a dictionary.
      • 
 
    
@@ -257,13 +257,13 @@ 
    besskge.dataset.KGDataset
 
    Parameters:
    
 
      
-
    • data (ndarray[Any, dtype[int32]]) – Numpy array of triples [head_id, relation_id, tail_id]. Shape
+
    • data (ndarray[Any, dtype[int32]]) – Numpy array of triples [head_id, relation_id, tail_id]. Shape
 (num_triples, 3).
      • 
-
    • split (Tuple[float, float, float]) – Tuple to set the train/validation/test split.
      • 
-
    • seed (int) – Random seed for the train/validation/test split.
      • 
-
    • entity_dict (Optional[List[str]]) – Optional entity labels by ID.
      • 
-
    • relation_dict (Optional[List[str]]) – Optional relation labels by ID.
      • 
-
    • type_offsets (Optional[Dict[str, int]]) – Offset of entity types
      • 
+
    • split (Tuple[float, float, float]) – Tuple to set the train/validation/test split.
      • 
+
    • seed (int) – Random seed for the train/validation/test split.
      • 
+
    • entity_dict (Optional[List[str]]) – Optional entity labels by ID.
      • 
+
    • relation_dict (Optional[List[str]]) – Optional relation labels by ID.
      • 
+
    • type_offsets (Optional[Dict[str, int]]) – Offset of entity types
      • 
 
    
 
    
 
    Return type:
    
@@ -277,7 +277,7 @@ 
    besskge.dataset.KGDataset
 
    
-property ht_types: Dict[str, ndarray[Any, dtype[int32]]] | None
    
+property ht_types: Dict[str, ndarray[Any, dtype[int32]]] | None
 
    If entities have types, type IDs of triples’ heads/tails;
 {part: int32[n_triple, {h_type, t_type}]}
    
 
    
@@ -288,7 +288,7 @@ 
    besskge.dataset.KGDatasetKGDataset object saved with KGDataset.save().
    
 
    
 
    Parameters:
    
-
    path (Path) – Path to saved KGDataset object.
    
+
    path (Path) – Path to saved KGDataset object.
    
 
    
 
    Return type:
    
 
    KGDataset
    
@@ -301,33 +301,33 @@ 
    besskge.dataset.KGDataset
 
    
-n_entity: int
    
+n_entity: int
 
    Number of entities (nodes) in the knowledge graph
    
 
    
 
 
    
 
    
-n_relation_type: int
    
+n_relation_type: int
 
    Number of relation types (edge labels) in the knowledge graph
    
 
    
 
 
    
 
    
-neg_heads: Optional[Dict[str, ndarray[Any, dtype[int32]]]] = None
    
+neg_heads: Optional[Dict[str, ndarray[Any, dtype[int32]]]] = None
 
    IDs of (possibly triple-specific) negative heads;
 {part: int32[n_triple or 1, n_neg_heads]}
    
 
    
 
 
    
 
    
-neg_tails: Optional[Dict[str, ndarray[Any, dtype[int32]]]] = None
    
+neg_tails: Optional[Dict[str, ndarray[Any, dtype[int32]]]] = None
 
    IDs of (possibly triple-specific) negative tails;
 {part: int32[n_triple or 1, n_neg_tails]}
    
 
    
 
 
    
 
    
-relation_dict: Optional[List[str]] = None
    
+relation_dict: Optional[List[str]] = None
 
    Relation type labels by ID; str[n_relation_type]
    
 
    
 
@@ -337,24 +337,24 @@ 
    besskge.dataset.KGDataset
 
    Parameters:
    
-
    out_file (Path) – Path to output file.
    
+
    out_file (Path) – Path to output file.
    
 
    
 
    Return type:
    
-
    None
    
+
    None
    
 
    
 
 
    
 
 
    
 
    
-triples: Dict[str, ndarray[Any, dtype[int32]]]
    
+triples: Dict[str, ndarray[Any, dtype[int32]]]
 
    List of (h_ID, r_ID, t_ID) triples, for each part of the dataset;
 {part: int32[n_triple, {h,r,t}]}
    
 
    
 
 
    
 
    
-type_offsets: Optional[Dict[str, int]] = None
    
+type_offsets: Optional[Dict[str, int]] = None
 
    If entities have types, IDs are assumed to be clustered by type;
 {entity_type: int}
    
 
    
diff --git a/generated/besskge.embedding.init_KGE_normal.html b/generated/besskge.embedding.init_KGE_normal.html
index bc7a81a..5c7883d 100644
--- a/generated/besskge.embedding.init_KGE_normal.html
+++ b/generated/besskge.embedding.init_KGE_normal.html
@@ -108,13 +108,13 @@ 
    besskge.embedding.init_KGE_normal
 
    Parameters:
    
 
      
-
    • embedding_table (Tensor) – Tensor of embedding parameters to initialize.
      • 
-
    • std (float) – Standard deviation. Default: 1.0.
      • 
-
    • divide_by_embedding_size (bool) – Rescale standard deviation by 1/row_size. Default: True.
      • 
+
    • embedding_table (Tensor) – Tensor of embedding parameters to initialize.
      • 
+
    • std (float) – Standard deviation. Default: 1.0.
      • 
+
    • divide_by_embedding_size (bool) – Rescale standard deviation by 1/row_size. Default: True.
      • 
 
    
 
    
 
    Return type:
    
-
    Tensor
    
+
    Tensor
    
 
    
 
    Returns:
    
 
    Initialized tensor.
    
diff --git a/generated/besskge.embedding.init_KGE_uniform.html b/generated/besskge.embedding.init_KGE_uniform.html
index 146248d..beab19d 100644
--- a/generated/besskge.embedding.init_KGE_uniform.html
+++ b/generated/besskge.embedding.init_KGE_uniform.html
@@ -108,13 +108,13 @@ 
    besskge.embedding.init_KGE_uniform
 
    Parameters:
    
 
      
-
    • embedding_table (Tensor) – Tensor of embedding parameters to initialize.
      • 
-
    • b (float) – Positive boundary of distribution support. Default: 1.0.
      • 
-
    • divide_by_embedding_size (bool) – Rescale distribution support by 1/row_size. Default: True.
      • 
+
    • embedding_table (Tensor) – Tensor of embedding parameters to initialize.
      • 
+
    • b (float) – Positive boundary of distribution support. Default: 1.0.
      • 
+
    • divide_by_embedding_size (bool) – Rescale distribution support by 1/row_size. Default: True.
      • 
 
    
 
    
 
    Return type:
    
-
    Tensor
    
+
    Tensor
    
 
    
 
    Returns:
    
 
    Initialized tensor.
    
diff --git a/generated/besskge.embedding.init_uniform_norm.html b/generated/besskge.embedding.init_uniform_norm.html
index 9641e88..dbc927c 100644
--- a/generated/besskge.embedding.init_uniform_norm.html
+++ b/generated/besskge.embedding.init_uniform_norm.html
@@ -108,10 +108,10 @@ 
    besskge.embedding.init_uniform_norm
 
    Parameters:
    
-
    embedding_table (Tensor) – Tensor of embedding parameters to initialize.
    
+
    embedding_table (Tensor) – Tensor of embedding parameters to initialize.
    
 
    
 
    Return type:
    
-
    Tensor
    
+
    Tensor
    
 
    
 
    Returns:
    
 
    Initialized tensor.
    
diff --git a/generated/besskge.embedding.initialize_entity_embedding.html b/generated/besskge.embedding.initialize_entity_embedding.html
index 945187b..2384db6 100644
--- a/generated/besskge.embedding.initialize_entity_embedding.html
+++ b/generated/besskge.embedding.initialize_entity_embedding.html
@@ -109,13 +109,13 @@ 
    besskge.embedding.initialize_entity_embeddingParameters:
 
      
 
    • sharding (Sharding) – Entity sharding.
      • 
-
    • initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Embedding table or list of initializing functions. If providing
+
    • 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.
+
    • 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.
@@ -123,7 +123,7 @@ 
      besskge.embedding.initialize_entity_embeddingReturn type:
-
      Parameter
      
+
      Parameter
      
 
      
 
      Returns:
      
 
      shape: (n_shard, max_ent_per_shard, row_size)
diff --git a/generated/besskge.embedding.initialize_relation_embedding.html b/generated/besskge.embedding.initialize_relation_embedding.html
index 2ee3ff9..e2acd99 100644
--- a/generated/besskge.embedding.initialize_relation_embedding.html
+++ b/generated/besskge.embedding.initialize_relation_embedding.html
@@ -108,15 +108,15 @@ 
      besskge.embedding.initialize_relation_embedding
 
      Parameters:
      
 
        
-
      • n_relation_type (int) – Number of relation types.
        • 
-
      • inverse_relations (bool) – If True, learn embeddings for inverse relations, in addition to direct ones.
+
      • n_relation_type (int) – Number of relation types.
        • 
+
      • inverse_relations (bool) – If True, learn embeddings for inverse relations, in addition to direct ones.
 Needs to be set to True when inverse triples are added to the dataset.
 Given a relation with ID i, its inverse is the one with
 ID i+n_relation_type.
        • 
-
      • initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Embedding table or list of initializing functions.
+
      • initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Embedding table or list of initializing functions.
 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 relation type.
+
      • row_size (Optional[List[int]]) – Number of parameters for each relation type.
 This needs to be a list, with the lengths of the different embedding tensors
 to allocate for each relation. Each embedding tensor, once allocated, is
 initialized with the corresponding entry of initializer.
@@ -124,7 +124,7 @@ 
        besskge.embedding.initialize_relation_embeddingReturn type:
-
        Parameter
        
+
        Parameter
        
 
        
 
        Returns:
        
 
        Relation embedding table.
        
diff --git a/generated/besskge.embedding.refactor_embedding_sharding.html b/generated/besskge.embedding.refactor_embedding_sharding.html
index ae5688d..a61c25a 100644
--- a/generated/besskge.embedding.refactor_embedding_sharding.html
+++ b/generated/besskge.embedding.refactor_embedding_sharding.html
@@ -109,14 +109,14 @@ 
        besskge.embedding.refactor_embedding_sharding
 
        Parameters:
        
 
          
-
        • entity_embedding (Parameter) – shape: (n_shard_old, max_ent_per_shard_old, row_size)
+
        • entity_embedding (Parameter) – shape: (n_shard_old, max_ent_per_shard_old, row_size)
 Entity embedding table sharded according to old_sharding.
          • 
 
        • old_sharding (Sharding) – The current entity sharding.
          • 
 
        • new_sharding (Sharding) – The new entity sharding.
          • 
 
        
 
        
 
        Return type:
        
-
        Parameter
        
+
        Parameter
        
 
        
 
        Returns:
        
 
        shape: (n_shard_new, max_ent_per_shard_new, row_size)
diff --git a/generated/besskge.loss.BaseLossFunction.html b/generated/besskge.loss.BaseLossFunction.html
index c16a533..f9933f1 100644
--- a/generated/besskge.loss.BaseLossFunction.html
+++ b/generated/besskge.loss.BaseLossFunction.html
@@ -115,16 +115,16 @@ 
        besskge.loss.BaseLossFunction
 
        Parameters:
        
 
          
-
        • positive_score (Tensor) – shape: (batch_size,)
+
        • positive_score (Tensor) – shape: (batch_size,)
 Scores of positive triples.
          • 
-
        • negative_score (Tensor) – shape: (batch_size, n_negative)
+
        • negative_score (Tensor) – shape: (batch_size, n_negative)
 Scores of negative triples.
          • 
-
        • triple_weight (Tensor) – shape: (batch_size,) or ()
+
        • triple_weight (Tensor) – shape: (batch_size,) or ()
 Weights of positive triples.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        The batch loss.
        
@@ -138,11 +138,11 @@ 
        besskge.loss.BaseLossFunction
 
        Parameters:
        
-
        negative_score (Tensor) – : (batch_size, n_negative)
+
        negative_score (Tensor) – : (batch_size, n_negative)
 Scores of negative samples.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, n_negative)
@@ -154,19 +154,19 @@ 
        besskge.loss.BaseLossFunction
 
        
-loss_scale: Tensor
        
+loss_scale: Tensor
 
        Loss scaling factor, might be needed when using FP16 weights
        
 
        
 
 
        
 
        
-negative_adversarial_sampling: bool
        
+negative_adversarial_sampling: bool
 
        Use self-adversarial weighting of negative samples.
        
 
        
 
 
        
 
        
-negative_adversarial_scale: Tensor
        
+negative_adversarial_scale: Tensor
 
        Reciprocal temperature of self-adversarial weighting
        
 
        
 
diff --git a/generated/besskge.loss.LogSigmoidLoss.html b/generated/besskge.loss.LogSigmoidLoss.html
index 5979c17..a5b41b9 100644
--- a/generated/besskge.loss.LogSigmoidLoss.html
+++ b/generated/besskge.loss.LogSigmoidLoss.html
@@ -108,10 +108,10 @@ 
        besskge.loss.LogSigmoidLoss
 
        Parameters:
        
 
        
 
        
 
@@ -122,16 +122,16 @@ 
        besskge.loss.LogSigmoidLoss
 
        Parameters:
        
 
          
-
        • positive_score (Tensor) – shape: (batch_size,)
+
        • positive_score (Tensor) – shape: (batch_size,)
 Scores of positive triples.
          • 
-
        • negative_score (Tensor) – shape: (batch_size, n_negative)
+
        • negative_score (Tensor) – shape: (batch_size, n_negative)
 Scores of negative triples.
          • 
-
        • triple_weight (Tensor) – shape: (batch_size,) or ()
+
        • triple_weight (Tensor) – shape: (batch_size,) or ()
 Weights of positive triples.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        The batch loss.
        
@@ -145,11 +145,11 @@ 
        besskge.loss.LogSigmoidLoss
 
        Parameters:
        
-
        negative_score (Tensor) – : (batch_size, n_negative)
+
        negative_score (Tensor) – : (batch_size, n_negative)
 Scores of negative samples.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, n_negative)
@@ -161,19 +161,19 @@ 
        besskge.loss.LogSigmoidLoss
 
        
-loss_scale: Tensor
        
+loss_scale: Tensor
 
        Loss scaling factor, might be needed when using FP16 weights
        
 
        
 
 
        
 
        
-negative_adversarial_sampling: bool
        
+negative_adversarial_sampling: bool
 
        Use self-adversarial weighting of negative samples.
        
 
        
 
 
        
 
        
-negative_adversarial_scale: Tensor
        
+negative_adversarial_scale: Tensor
 
        Reciprocal temperature of self-adversarial weighting
        
 
        
 
diff --git a/generated/besskge.loss.MarginBasedLossFunction.html b/generated/besskge.loss.MarginBasedLossFunction.html
index 711a137..5572a81 100644
--- a/generated/besskge.loss.MarginBasedLossFunction.html
+++ b/generated/besskge.loss.MarginBasedLossFunction.html
@@ -108,10 +108,10 @@ 
        besskge.loss.MarginBasedLossFunction
 
        Parameters:
        
 
        
 
        
 
@@ -122,16 +122,16 @@ 
        besskge.loss.MarginBasedLossFunction
 
        Parameters:
        
 
          
-
        • positive_score (Tensor) – shape: (batch_size,)
+
        • positive_score (Tensor) – shape: (batch_size,)
 Scores of positive triples.
          • 
-
        • negative_score (Tensor) – shape: (batch_size, n_negative)
+
        • negative_score (Tensor) – shape: (batch_size, n_negative)
 Scores of negative triples.
          • 
-
        • triple_weight (Tensor) – shape: (batch_size,) or ()
+
        • triple_weight (Tensor) – shape: (batch_size,) or ()
 Weights of positive triples.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        The batch loss.
        
@@ -145,11 +145,11 @@ 
        besskge.loss.MarginBasedLossFunction
 
        Parameters:
        
-
        negative_score (Tensor) – : (batch_size, n_negative)
+
        negative_score (Tensor) – : (batch_size, n_negative)
 Scores of negative samples.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, n_negative)
@@ -161,19 +161,19 @@ 
        besskge.loss.MarginBasedLossFunction
 
        
-loss_scale: Tensor
        
+loss_scale: Tensor
 
        Loss scaling factor, might be needed when using FP16 weights
        
 
        
 
 
        
 
        
-negative_adversarial_sampling: bool
        
+negative_adversarial_sampling: bool
 
        Use self-adversarial weighting of negative samples.
        
 
        
 
 
        
 
        
-negative_adversarial_scale: Tensor
        
+negative_adversarial_scale: Tensor
 
        Reciprocal temperature of self-adversarial weighting
        
 
        
 
diff --git a/generated/besskge.loss.MarginRankingLoss.html b/generated/besskge.loss.MarginRankingLoss.html
index 984f57a..d5f8868 100644
--- a/generated/besskge.loss.MarginRankingLoss.html
+++ b/generated/besskge.loss.MarginRankingLoss.html
@@ -108,11 +108,11 @@ 
        besskge.loss.MarginRankingLoss
 
        Parameters:
        
 
        
 
        
 
@@ -123,16 +123,16 @@ 
        besskge.loss.MarginRankingLoss
 
        Parameters:
        
 
          
-
        • positive_score (Tensor) – shape: (batch_size,)
+
        • positive_score (Tensor) – shape: (batch_size,)
 Scores of positive triples.
          • 
-
        • negative_score (Tensor) – shape: (batch_size, n_negative)
+
        • negative_score (Tensor) – shape: (batch_size, n_negative)
 Scores of negative triples.
          • 
-
        • triple_weight (Tensor) – shape: (batch_size,) or ()
+
        • triple_weight (Tensor) – shape: (batch_size,) or ()
 Weights of positive triples.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        The batch loss.
        
@@ -146,11 +146,11 @@ 
        besskge.loss.MarginRankingLoss
 
        Parameters:
        
-
        negative_score (Tensor) – : (batch_size, n_negative)
+
        negative_score (Tensor) – : (batch_size, n_negative)
 Scores of negative samples.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, n_negative)
@@ -162,19 +162,19 @@ 
        besskge.loss.MarginRankingLoss
 
        
-loss_scale: Tensor
        
+loss_scale: Tensor
 
        Loss scaling factor, might be needed when using FP16 weights
        
 
        
 
 
        
 
        
-negative_adversarial_sampling: bool
        
+negative_adversarial_sampling: bool
 
        Use self-adversarial weighting of negative samples.
        
 
        
 
 
        
 
        
-negative_adversarial_scale: Tensor
        
+negative_adversarial_scale: Tensor
 
        Reciprocal temperature of self-adversarial weighting
        
 
        
 
diff --git a/generated/besskge.loss.SampledSoftmaxCrossEntropyLoss.html b/generated/besskge.loss.SampledSoftmaxCrossEntropyLoss.html
index f523ef7..abbd947 100644
--- a/generated/besskge.loss.SampledSoftmaxCrossEntropyLoss.html
+++ b/generated/besskge.loss.SampledSoftmaxCrossEntropyLoss.html
@@ -109,8 +109,8 @@ 
        besskge.loss.SampledSoftmaxCrossEntropyLoss
 
        Parameters:
        
 
        
 
        
 
@@ -121,16 +121,16 @@ 
        besskge.loss.SampledSoftmaxCrossEntropyLoss
 
        Parameters:
        
 
          
-
        • positive_score (Tensor) – shape: (batch_size,)
+
        • positive_score (Tensor) – shape: (batch_size,)
 Scores of positive triples.
          • 
-
        • negative_score (Tensor) – shape: (batch_size, n_negative)
+
        • negative_score (Tensor) – shape: (batch_size, n_negative)
 Scores of negative triples.
          • 
-
        • triple_weight (Tensor) – shape: (batch_size,) or ()
+
        • triple_weight (Tensor) – shape: (batch_size,) or ()
 Weights of positive triples.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        The batch loss.
        
@@ -144,11 +144,11 @@ 
        besskge.loss.SampledSoftmaxCrossEntropyLoss
 
        Parameters:
        
-
        negative_score (Tensor) – : (batch_size, n_negative)
+
        negative_score (Tensor) – : (batch_size, n_negative)
 Scores of negative samples.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, n_negative)
@@ -160,19 +160,19 @@ 
        besskge.loss.SampledSoftmaxCrossEntropyLoss
 
        
-loss_scale: Tensor
        
+loss_scale: Tensor
 
        Loss scaling factor, might be needed when using FP16 weights
        
 
        
 
 
        
 
        
-negative_adversarial_sampling: bool
        
+negative_adversarial_sampling: bool
 
        Use self-adversarial weighting of negative samples.
        
 
        
 
 
        
 
        
-negative_adversarial_scale: Tensor
        
+negative_adversarial_scale: Tensor
 
        Reciprocal temperature of self-adversarial weighting
        
 
        
 
diff --git a/generated/besskge.metric.Evaluation.html b/generated/besskge.metric.Evaluation.html
index 3da7b45..33b91d5 100644
--- a/generated/besskge.metric.Evaluation.html
+++ b/generated/besskge.metric.Evaluation.html
@@ -108,15 +108,15 @@ 
        besskge.metric.Evaluation
 
        Parameters:
        
 
          
-
        • metric_list (List[str]) – List of metrics to compute. Currently supports “mrr” and “hits@K”.
          • 
-
        • mode (str) – Mode used for metrics. Can be “optimistic”, “pessimistic”
+
        • metric_list (List[str]) – List of metrics to compute. Currently supports “mrr” and “hits@K”.
          • 
+
        • mode (str) – Mode used for metrics. Can be “optimistic”, “pessimistic”
 or “average”. Default: “average”.
          • 
-
        • worst_rank_infty (bool) – If True, assign a prediction rank of infinity as the worst possible
+
        • worst_rank_infty (bool) – If True, assign a prediction rank of infinity as the worst possible
 rank. If False, assign a prediction rank of n_negative + 1 as the
 worst possible rank. Default: False.
          • 
-
        • reduction (str) – Method to use to reduce metrics along the batch dimension.
+
        • reduction (str) – Method to use to reduce metrics along the batch dimension.
 Currently supports “none” (no reduction) and “sum”.
          • 
-
        • return_ranks (bool) – If True, returns prediction ranks alongside metrics.
          • 
+
        • return_ranks (bool) – If True, returns prediction ranks alongside metrics.
          • 
 
        
 
        
 
@@ -128,15 +128,15 @@ 
        besskge.metric.Evaluation
 
        Parameters:
        
 
          
-
        • batch_rank (Tensor) – shape: (batch_size,)
+
        • batch_rank (Tensor) – shape: (batch_size,)
 Prediction rank for each element in the batch.
          • 
-
        • triple_mask (Optional[Tensor]) – shape: (batch_size,)
+
        • triple_mask (Optional[Tensor]) – shape: (batch_size,)
 Boolean mask. If provided, all metrics for the elements where
 ~triple_mask are set to 0.0.
          • 
 
        
 
        
 
        Return type:
        
-
        Dict[str, Tensor]
        
+
        Dict[str, Tensor]
        
 
        
 
        Returns:
        
 
        The dictionary of (reduced) batch metrics.
        
@@ -152,16 +152,16 @@ 
        besskge.metric.Evaluation
 
        Parameters:
        
 
          
-
        • ground_truth (Tensor) – shape: (batch_size,)
+
        • ground_truth (Tensor) – shape: (batch_size,)
 Indices of ground truth entities for each query.
          • 
-
        • candidate_indices (Tensor) – shape: (batch_size, n_candidates)
+
        • candidate_indices (Tensor) – shape: (batch_size, n_candidates)
 Indices of top n_candidates predicted entities,
 ordered by decreasing likelihood. The indices
 on each row are assumed to be distinct.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        The rank of the ground truth among the predictions.
        
@@ -178,14 +178,14 @@ 
        besskge.metric.Evaluation
 
        Parameters:
        
 
          
-
        • pos_score (Tensor) – shape: (batch_size,)
+
        • pos_score (Tensor) – shape: (batch_size,)
 Scores of positive triples.
          • 
-
        • candidate_score (Tensor) – shape: (batch_size, n_candidate)
+
        • candidate_score (Tensor) – shape: (batch_size, n_candidate)
 Scores of candidate triples.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        The rank of the positive score among the ordered scores
@@ -203,14 +203,14 @@ 
        besskge.metric.Evaluation
 
        Parameters:
        
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (1, n_metrics, batch_size)
diff --git a/generated/besskge.metric.HitsAtK.html b/generated/besskge.metric.HitsAtK.html
index 81b3257..fefbed1 100644
--- a/generated/besskge.metric.HitsAtK.html
+++ b/generated/besskge.metric.HitsAtK.html
@@ -108,7 +108,7 @@ 
        besskge.metric.HitsAtK
 
        
 
        Parameters:
        
-
        k (int) – Maximum acceptable rank.
        
+
        k (int) – Maximum acceptable rank.
        
 
        
 
        
 
        
diff --git a/generated/besskge.negative_sampler.PlaceholderNegativeSampler.html b/generated/besskge.negative_sampler.PlaceholderNegativeSampler.html
index e1b5ce1..0706d58 100644
--- a/generated/besskge.negative_sampler.PlaceholderNegativeSampler.html
+++ b/generated/besskge.negative_sampler.PlaceholderNegativeSampler.html
@@ -110,26 +110,26 @@ 
        besskge.negative_sampler.PlaceholderNegativeSampler
 
        Parameters:
        
 
        
 
        
 
 
        
 
        
-corruption_scheme: str
        
+corruption_scheme: str
 
        Which entity to corrupt; “h”, “t”, “ht”
        
 
        
 
 
        
 
        
-flat_negative_format: bool
        
+flat_negative_format: bool
 
        Sample negatives per triple partition, instead of per triple
        
 
        
 
 
        
 
        
-local_sampling: bool
        
+local_sampling: bool
 
        Sample negatives only from processing device
        
 
        
 
diff --git a/generated/besskge.negative_sampler.RandomShardedNegativeSampler.html b/generated/besskge.negative_sampler.RandomShardedNegativeSampler.html
index 82ebd5a..1de1dd6 100644
--- a/generated/besskge.negative_sampler.RandomShardedNegativeSampler.html
+++ b/generated/besskge.negative_sampler.RandomShardedNegativeSampler.html
@@ -108,17 +108,17 @@ 
        besskge.negative_sampler.RandomShardedNegativeSampler
 
        Parameters:
        
 
          
-
        • n_negative (int) – Number of negative samples per shard-pair
+
        • n_negative (int) – Number of negative samples per shard-pair
 (if flat_negative_format) or per triple.
          • 
 
        • sharding (Sharding) – Sharding of entities.
          • 
-
        • seed (int) – Seed of RNG.
          • 
-
        • corruption_scheme (str) – “h”: corrupt head entities;
+
        • seed (int) – Seed of RNG.
          • 
+
        • corruption_scheme (str) – “h”: corrupt head entities;
 “t”: corrupt tail entities;
 “ht”: corrupt head entities for the first half of each triple
 partition, tail entities for the second half.
          • 
-
        • local_sampling (bool) – If True, sample negative entities only from the shard where the
+
        • local_sampling (bool) – If True, sample negative entities only from the shard where the
 triple is processed.
          • 
-
        • flat_negative_format (bool) – If True, sample n_negative negative entities for each
+
        • flat_negative_format (bool) – If True, sample n_negative negative entities for each
 shard-pair, instead of each triple. If True, requires use of
 negative sample sharing. Default: False.
          • 
 
        
@@ -126,19 +126,19 @@ 
        besskge.negative_sampler.RandomShardedNegativeSampler
 
        
 
        
-corruption_scheme: str
        
+corruption_scheme: str
 
        Which entity to corrupt; “h”, “t”, “ht”
        
 
        
 
 
        
 
        
-flat_negative_format: bool
        
+flat_negative_format: bool
 
        Sample negatives per triple partition, instead of per triple
        
 
        
 
 
        
 
        
-local_sampling: bool
        
+local_sampling: bool
 
        Sample negatives only from processing device
        
 
        
 
diff --git a/generated/besskge.negative_sampler.ShardedNegativeSampler.html b/generated/besskge.negative_sampler.ShardedNegativeSampler.html
index 3e2398a..ee75865 100644
--- a/generated/besskge.negative_sampler.ShardedNegativeSampler.html
+++ b/generated/besskge.negative_sampler.ShardedNegativeSampler.html
@@ -106,19 +106,19 @@ 
        besskge.negative_sampler.ShardedNegativeSampler
 
        
-corruption_scheme: str
        
+corruption_scheme: str
 
        Which entity to corrupt; “h”, “t”, “ht”
        
 
        
 
 
        
 
        
-flat_negative_format: bool
        
+flat_negative_format: bool
 
        Sample negatives per triple partition, instead of per triple
        
 
        
 
 
        
 
        
-local_sampling: bool
        
+local_sampling: bool
 
        Sample negatives only from processing device
        
 
        
 
diff --git a/generated/besskge.negative_sampler.TripleBasedShardedNegativeSampler.html b/generated/besskge.negative_sampler.TripleBasedShardedNegativeSampler.html
index 4b80057..ebde14b 100644
--- a/generated/besskge.negative_sampler.TripleBasedShardedNegativeSampler.html
+++ b/generated/besskge.negative_sampler.TripleBasedShardedNegativeSampler.html
@@ -108,20 +108,20 @@ 
        besskge.negative_sampler.TripleBasedShardedNegativeSampler
 
        Parameters:
        
 
          
-
        • negative_heads (Optional[ndarray[Any, dtype[int32]]]) – shape: (N, n_negative)
+
        • negative_heads (Optional[ndarray[Any, dtype[int32]]]) – shape: (N, n_negative)
 Global entity IDs of negative heads, specific
 for each triple (N=n_triple) or for all of them (N=1).
          • 
-
        • negative_tails (Optional[ndarray[Any, dtype[int32]]]) – shape: (N, n_negative)
+
        • negative_tails (Optional[ndarray[Any, dtype[int32]]]) – shape: (N, n_negative)
 Global entity IDs of negative tails, specific
 for each triple (N=n_triple) or for all of them (N=1).
          • 
 
        • sharding (Sharding) – see RandomShardedNegativeSampler.__init__()
          • 
-
        • corruption_scheme (str) – see RandomShardedNegativeSampler.__init__()
          • 
-
        • seed (int) – see RandomShardedNegativeSampler.__init__()
          • 
-
        • mask_on_gather (bool) – If True, shape the negative mask to be applied on the device where
+
        • corruption_scheme (str) – see RandomShardedNegativeSampler.__init__()
          • 
+
        • seed (int) – see RandomShardedNegativeSampler.__init__()
          • 
+
        • mask_on_gather (bool) – If True, shape the negative mask to be applied on the device where
 negative entities are gathered, instead of the one where they
 are scored. Set to True only when using
 besskge.bess.TopKQueryBessKGE. Default: False.
          • 
-
        • return_sort_idx (bool) – If True, return for each triple in the batch the sorting indices
+
        • return_sort_idx (bool) – If True, return for each triple in the batch the sorting indices
 to recover the same ordering of negatives as in
 negative_heads, negative_tails. Default: False.
          • 
 
        
@@ -129,19 +129,19 @@ 
        besskge.negative_sampler.TripleBasedShardedNegativeSampler
 
        
-corruption_scheme: str
        
+corruption_scheme: str
 
        Which entity to corrupt; “h”, “t”, “ht”
        
 
        
 
 
        
 
        
-flat_negative_format: bool
        
+flat_negative_format: bool
 
        Sample negatives per triple partition, instead of per triple
        
 
        
 
 
        
 
        
-local_sampling: bool
        
+local_sampling: bool
 
        Sample negatives only from processing device
        
 
        
 
@@ -152,12 +152,12 @@ 
        besskge.negative_sampler.TripleBasedShardedNegativeSampler
 
        Parameters:
        
 
          
-
        • negatives (ndarray[Any, dtype[int32]]) – shape: (N, n_negative)
+
        • negatives (ndarray[Any, dtype[int32]]) – shape: (N, n_negative)
 Negative entities, each row already sorted in shard order
 (N = 1, n_triple).
          • 
-
        • shard_counts (ndarray[Any, dtype[int64]]) – shape: (N, n_shard)
+
        • shard_counts (ndarray[Any, dtype[int64]]) – shape: (N, n_shard)
 Number of negatives per shard.
          • 
-
        • padded_shard_length (int) – The size to which each shard list is to be padded.
          • 
+
        • padded_shard_length (int) – The size to which each shard list is to be padded.
          • 
 
        
 
        
 
        Return padded_negatives:
        
@@ -169,7 +169,7 @@ 
        besskge.negative_sampler.TripleBasedShardedNegativeSamplerpadded_negatives.view(N,-1).
        
 
        
 
        Return type:
        
-
        Tuple[ndarray[Any, dtype[int32]], ndarray[Any, dtype[bool_]]]
        
+
        Tuple[ndarray[Any, dtype[int32]], ndarray[Any, dtype[bool_]]]
        
 
        
 
 
        
@@ -186,7 +186,7 @@ 
        besskge.negative_sampler.TripleBasedShardedNegativeSampler
 
        Parameters:
        
-
        negatives (ndarray[Any, dtype[int32]]) – shape: (N, n_negatives)
+
        negatives (ndarray[Any, dtype[int32]]) – shape: (N, n_negatives)
 Negative entities to shard (N = 1, n_triple).
        
 
        
 
        Return shard_neg_counts:
        
@@ -198,7 +198,7 @@ 
        besskge.negative_sampler.TripleBasedShardedNegativeSamplerReturn type:
-
        Tuple[ndarray[Any, dtype[int64]], ndarray[Any, dtype[int32]]]
        
+
        Tuple[ndarray[Any, dtype[int64]], ndarray[Any, dtype[int32]]]
        
 
        
 
 
        
diff --git a/generated/besskge.negative_sampler.TypeBasedShardedNegativeSampler.html b/generated/besskge.negative_sampler.TypeBasedShardedNegativeSampler.html
index c226711..9e4cdaa 100644
--- a/generated/besskge.negative_sampler.TypeBasedShardedNegativeSampler.html
+++ b/generated/besskge.negative_sampler.TypeBasedShardedNegativeSampler.html
@@ -108,31 +108,31 @@ 
        besskge.negative_sampler.TypeBasedShardedNegativeSampler
 
        Parameters:
        
 
          
-
        • triple_types (ndarray[Any, dtype[int32]]) – shape: (n_triple, 2)
+
        • triple_types (ndarray[Any, dtype[int32]]) – shape: (n_triple, 2)
 Type IDs of head and tail entities for all triples.
          • 
-
        • n_negative (int) – see RandomShardedNegativeSampler.__init__()
          • 
+
        • n_negative (int) – see RandomShardedNegativeSampler.__init__()
          • 
 
        • sharding (Sharding) – see RandomShardedNegativeSampler.__init__()
          • 
-
        • corruption_scheme (str) – see RandomShardedNegativeSampler.__init__()
          • 
-
        • local_sampling (bool) – see RandomShardedNegativeSampler.__init__()
          • 
-
        • seed (int) – see RandomShardedNegativeSampler.__init__()
          • 
+
        • corruption_scheme (str) – see RandomShardedNegativeSampler.__init__()
          • 
+
        • local_sampling (bool) – see RandomShardedNegativeSampler.__init__()
          • 
+
        • seed (int) – see RandomShardedNegativeSampler.__init__()
          • 
 
        
 
        
 
 
        
 
        
-corruption_scheme: str
        
+corruption_scheme: str
 
        Which entity to corrupt; “h”, “t”, “ht”
        
 
        
 
 
        
 
        
-flat_negative_format: bool
        
+flat_negative_format: bool
 
        Sample negatives per triple partition, instead of per triple
        
 
        
 
 
        
 
        
-local_sampling: bool
        
+local_sampling: bool
 
        Sample negatives only from processing device
        
 
        
 
diff --git a/generated/besskge.scoring.BaseScoreFunction.html b/generated/besskge.scoring.BaseScoreFunction.html
index cae8c21..3f6d410 100644
--- a/generated/besskge.scoring.BaseScoreFunction.html
+++ b/generated/besskge.scoring.BaseScoreFunction.html
@@ -116,7 +116,7 @@ 
        besskge.scoring.BaseScoreFunction
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -126,13 +126,13 @@ 
        besskge.scoring.BaseScoreFunctionBaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -140,13 +140,13 @@ 
        besskge.scoring.BaseScoreFunction
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
 
        
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -157,16 +157,16 @@ 
        besskge.scoring.BaseScoreFunction
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -184,16 +184,16 @@ 
        besskge.scoring.BaseScoreFunction
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -211,16 +211,16 @@ 
        besskge.scoring.BaseScoreFunction
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -244,7 +244,7 @@ 
        besskge.scoring.BaseScoreFunction
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.BoxE.html b/generated/besskge.scoring.BoxE.html
index 88982aa..28fb804 100644
--- a/generated/besskge.scoring.BoxE.html
+++ b/generated/besskge.scoring.BoxE.html
@@ -115,25 +115,25 @@ 
        besskge.scoring.BoxE
 
        Parameters:
        
 
          
-
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
-
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
 
        • sharding (Sharding) – Entity sharding.
          • 
-
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
-
        • embedding_size (int) – Size of final entity embeddings.
          • 
-
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
-
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization functions or table for relation embeddings.
+
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
+
        • embedding_size (int) – Size of final entity embeddings.
          • 
+
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
+
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization functions or table for relation embeddings.
 If not passing a table, two functions are needed: the initializer
 for head/tail box centers and the initializer for
 (scalar) head/tail box sizes.
          • 
-
        • apply_tanh (bool) – If True, bound relation box sizes and bumped entity
+
        • apply_tanh (bool) – If True, bound relation box sizes and bumped entity
 representations with tanh. Default: True.
          • 
-
        • dist_func_per_dim (bool) – If True, instead of selecting between the two BoxE distance
+
        • dist_func_per_dim (bool) – If True, instead of selecting between the two BoxE distance
 functions based on whether the bumped representation is inside
 or outside the relation box, make the choice separately for
 each dimension of the embedding space. Default: True.
          • 
-
        • eps (float) – Softening parameter for geometric normalization of box widths.
+
        • eps (float) – Softening parameter for geometric normalization of box widths.
 Default: 1e-6.
          • 
-
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
+
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
 
        
 
        
 
@@ -144,25 +144,25 @@ 
        besskge.scoring.BoxE
 
        Parameters:
        
 
          
-
        • bumped_ht (Tensor) – shape: (batch_size, 2, emb_size)
+
        • bumped_ht (Tensor) – shape: (batch_size, 2, emb_size)
 or (batch_size, n_negative, 2, emb_size)
 Bumped h/t entity embeddings (heads: bumped_ht[…,0,:],
 tails: bumped_ht[…,1,:]).
          • 
-
        • center_ht (Tensor) – shape: (batch_size, 2, emb_size)
+
        • center_ht (Tensor) – shape: (batch_size, 2, emb_size)
 or (batch_size, 1, 2, emb_size)
 Centers of h/t relation boxes
 (heads: center_ht[…,0,:], tails: center_ht[…,1,:]).
          • 
-
        • width_ht (Tensor) – shape: (batch_size, 2, emb_size) or (batch_size, 1, 2, emb_size)
+
        • width_ht (Tensor) – shape: (batch_size, 2, emb_size) or (batch_size, 1, 2, emb_size)
 Widths of h/t relation boxes, before normalization
 (heads: width_ht[…,0,:], tails: width_ht[…,1,:]).
          • 
-
        • box_size (Tensor) – shape: (batch_size, 2) or (batch_size, 1, 2)
+
        • box_size (Tensor) – shape: (batch_size, 2) or (batch_size, 1, 2)
 Parameter controlling the size of the (normalized)
 h/t relation boxes
 (heads: box_size[…,0], tails: box_size[…,1]).
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: shape: (batch_size,) or (batch_size, n_negative)
@@ -180,14 +180,14 @@ 
        besskge.scoring.BoxE
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -199,7 +199,7 @@ 
        besskge.scoring.BoxE
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -209,13 +209,13 @@ 
        besskge.scoring.BoxE
        see BaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -223,7 +223,7 @@ 
        besskge.scoring.BoxE
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -233,11 +233,11 @@ 
        besskge.scoring.BoxE
        p-norm reduction along embedding dimension.
        
 
        
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -248,7 +248,7 @@ 
        besskge.scoring.BoxE
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -259,16 +259,16 @@ 
        besskge.scoring.BoxE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -286,16 +286,16 @@ 
        besskge.scoring.BoxE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -313,16 +313,16 @@ 
        besskge.scoring.BoxE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -346,7 +346,7 @@ 
        besskge.scoring.BoxE
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.ComplEx.html b/generated/besskge.scoring.ComplEx.html
index f6a95d5..a11d72e 100644
--- a/generated/besskge.scoring.ComplEx.html
+++ b/generated/besskge.scoring.ComplEx.html
@@ -115,13 +115,13 @@ 
        besskge.scoring.ComplEx
 
        Parameters:
        
 
          
-
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
 
        • sharding (Sharding) – Entity sharding.
          • 
-
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
-
        • embedding_size (int) – Complex size of entity and relation embeddings.
          • 
-
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
-
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
-
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
+
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
+
        • embedding_size (int) – Complex size of entity and relation embeddings.
          • 
+
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
+
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
+
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
 
        
 
        
 
@@ -134,14 +134,14 @@ 
        besskge.scoring.ComplEx
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -153,7 +153,7 @@ 
        besskge.scoring.ComplEx
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -163,13 +163,13 @@ 
        besskge.scoring.ComplEx
        see BaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -177,7 +177,7 @@ 
        besskge.scoring.ComplEx
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -187,11 +187,11 @@ 
        besskge.scoring.ComplEx
        Sum reduction along the embedding dimension.
        
 
        
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -202,7 +202,7 @@ 
        besskge.scoring.ComplEx
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -213,16 +213,16 @@ 
        besskge.scoring.ComplEx
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -240,16 +240,16 @@ 
        besskge.scoring.ComplEx
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -267,16 +267,16 @@ 
        besskge.scoring.ComplEx
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -300,7 +300,7 @@ 
        besskge.scoring.ComplEx
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.DistMult.html b/generated/besskge.scoring.DistMult.html
index 24433df..79d9b42 100644
--- a/generated/besskge.scoring.DistMult.html
+++ b/generated/besskge.scoring.DistMult.html
@@ -115,13 +115,13 @@ 
        besskge.scoring.DistMult
 
        Parameters:
        
 
          
-
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
 
        • sharding (Sharding) – Entity sharding.
          • 
-
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
-
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
-
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
-
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
-
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
+
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
+
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
+
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
+
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
+
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
 
        
 
        
 
@@ -134,14 +134,14 @@ 
        besskge.scoring.DistMult
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -153,7 +153,7 @@ 
        besskge.scoring.DistMult
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -163,13 +163,13 @@ 
        besskge.scoring.DistMultBaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -177,7 +177,7 @@ 
        besskge.scoring.DistMult
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -187,11 +187,11 @@ 
        besskge.scoring.DistMult
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -202,7 +202,7 @@ 
        besskge.scoring.DistMult
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -213,16 +213,16 @@ 
        besskge.scoring.DistMult
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -240,16 +240,16 @@ 
        besskge.scoring.DistMult
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -267,16 +267,16 @@ 
        besskge.scoring.DistMult
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -300,7 +300,7 @@ 
        besskge.scoring.DistMult
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.DistanceBasedScoreFunction.html b/generated/besskge.scoring.DistanceBasedScoreFunction.html
index ec249c3..99b3239 100644
--- a/generated/besskge.scoring.DistanceBasedScoreFunction.html
+++ b/generated/besskge.scoring.DistanceBasedScoreFunction.html
@@ -115,8 +115,8 @@ 
        besskge.scoring.DistanceBasedScoreFunction
 
        Parameters:
        
 
        
 
        
 
@@ -129,14 +129,14 @@ 
        besskge.scoring.DistanceBasedScoreFunction
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -148,7 +148,7 @@ 
        besskge.scoring.DistanceBasedScoreFunction
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -158,13 +158,13 @@ 
        besskge.scoring.DistanceBasedScoreFunctionBaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -172,7 +172,7 @@ 
        besskge.scoring.DistanceBasedScoreFunction
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -182,11 +182,11 @@ 
        besskge.scoring.DistanceBasedScoreFunction
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -197,7 +197,7 @@ 
        besskge.scoring.DistanceBasedScoreFunction
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -208,16 +208,16 @@ 
        besskge.scoring.DistanceBasedScoreFunction
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -235,16 +235,16 @@ 
        besskge.scoring.DistanceBasedScoreFunction
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -262,16 +262,16 @@ 
        besskge.scoring.DistanceBasedScoreFunction
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -295,7 +295,7 @@ 
        besskge.scoring.DistanceBasedScoreFunction
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.InterHT.html b/generated/besskge.scoring.InterHT.html
index a4575d5..ee1cc46 100644
--- a/generated/besskge.scoring.InterHT.html
+++ b/generated/besskge.scoring.InterHT.html
@@ -115,17 +115,17 @@ 
        besskge.scoring.InterHT
 
        Parameters:
        
 
          
-
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
-
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
 
        • sharding (Sharding) – Entity sharding.
          • 
-
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
-
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
-
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
-
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
-
        • normalize_entities (bool) – If True, L2-normalize embeddings of head and tail entities as well as
+
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
+
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
+
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
+
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
+
        • normalize_entities (bool) – If True, L2-normalize embeddings of head and tail entities as well as
 auxiliary head and tail entities before multiplying. Default: True.
          • 
-
        • offset (float) – Offset applied to auxiliary entity embeddings. Default: 1.0.
          • 
-
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
+
        • offset (float) – Offset applied to auxiliary entity embeddings. Default: 1.0.
          • 
+
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
 
        
 
        
 
@@ -138,14 +138,14 @@ 
        besskge.scoring.InterHT
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -157,7 +157,7 @@ 
        besskge.scoring.InterHT
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -167,13 +167,13 @@ 
        besskge.scoring.InterHT
        see BaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -181,7 +181,7 @@ 
        besskge.scoring.InterHT
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -191,11 +191,11 @@ 
        besskge.scoring.InterHT
        p-norm reduction along embedding dimension.
        
 
        
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -206,7 +206,7 @@ 
        besskge.scoring.InterHT
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -217,16 +217,16 @@ 
        besskge.scoring.InterHT
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -244,16 +244,16 @@ 
        besskge.scoring.InterHT
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -271,16 +271,16 @@ 
        besskge.scoring.InterHT
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -304,7 +304,7 @@ 
        besskge.scoring.InterHT
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.MatrixDecompositionScoreFunction.html b/generated/besskge.scoring.MatrixDecompositionScoreFunction.html
index 7ea076e..ceb05fc 100644
--- a/generated/besskge.scoring.MatrixDecompositionScoreFunction.html
+++ b/generated/besskge.scoring.MatrixDecompositionScoreFunction.html
@@ -114,7 +114,7 @@ 
        besskge.scoring.MatrixDecompositionScoreFunction
 
        Parameters:
        
-
        negative_sample_sharing (bool) – see BaseScoreFunction
        
+
        negative_sample_sharing (bool) – see BaseScoreFunction
        
 
        
 
 
        
@@ -126,14 +126,14 @@ 
        besskge.scoring.MatrixDecompositionScoreFunction
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -145,7 +145,7 @@ 
        besskge.scoring.MatrixDecompositionScoreFunction
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -155,13 +155,13 @@ 
        besskge.scoring.MatrixDecompositionScoreFunctionBaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -169,7 +169,7 @@ 
        besskge.scoring.MatrixDecompositionScoreFunction
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -179,11 +179,11 @@ 
        besskge.scoring.MatrixDecompositionScoreFunction
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -194,7 +194,7 @@ 
        besskge.scoring.MatrixDecompositionScoreFunction
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -205,16 +205,16 @@ 
        besskge.scoring.MatrixDecompositionScoreFunction
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -232,16 +232,16 @@ 
        besskge.scoring.MatrixDecompositionScoreFunction
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -259,16 +259,16 @@ 
        besskge.scoring.MatrixDecompositionScoreFunction
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -292,7 +292,7 @@ 
        besskge.scoring.MatrixDecompositionScoreFunction
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.PairRE.html b/generated/besskge.scoring.PairRE.html
index 8723f6b..9918c67 100644
--- a/generated/besskge.scoring.PairRE.html
+++ b/generated/besskge.scoring.PairRE.html
@@ -115,16 +115,16 @@ 
        besskge.scoring.PairRE
 
        Parameters:
        
 
          
-
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
-
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
 
        • sharding (Sharding) – Entity sharding.
          • 
-
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
-
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
-
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
-
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
-
        • normalize_entities (bool) – If True, L2-normalize head and tail entity embeddings before projecting,
+
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
+
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
+
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
+
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
+
        • normalize_entities (bool) – If True, L2-normalize head and tail entity embeddings before projecting,
 as in [CHWC21]. Default: True.
          • 
-
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
+
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
 
        
 
        
 
@@ -137,14 +137,14 @@ 
        besskge.scoring.PairRE
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -156,7 +156,7 @@ 
        besskge.scoring.PairRE
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -166,13 +166,13 @@ 
        besskge.scoring.PairRE
        see BaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -180,7 +180,7 @@ 
        besskge.scoring.PairRE
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -190,11 +190,11 @@ 
        besskge.scoring.PairRE
        p-norm reduction along embedding dimension.
        
 
        
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -205,7 +205,7 @@ 
        besskge.scoring.PairRE
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -216,16 +216,16 @@ 
        besskge.scoring.PairRE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -243,16 +243,16 @@ 
        besskge.scoring.PairRE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -270,16 +270,16 @@ 
        besskge.scoring.PairRE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -303,7 +303,7 @@ 
        besskge.scoring.PairRE
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.RotatE.html b/generated/besskge.scoring.RotatE.html
index a532e49..40cdfb6 100644
--- a/generated/besskge.scoring.RotatE.html
+++ b/generated/besskge.scoring.RotatE.html
@@ -115,15 +115,15 @@ 
        besskge.scoring.RotatE
 
        Parameters:
        
 
          
-
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
-
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
 
        • sharding (Sharding) – Entity sharding.
          • 
-
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
-
        • embedding_size (int) – Complex size of entity embeddings (and real size of
+
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
+
        • embedding_size (int) – Complex size of entity embeddings (and real size of
 relation embeddings).
          • 
-
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
-
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
-
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
+
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
+
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
+
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
 
        
 
        
 
@@ -136,14 +136,14 @@ 
        besskge.scoring.RotatE
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -155,7 +155,7 @@ 
        besskge.scoring.RotatE
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -165,13 +165,13 @@ 
        besskge.scoring.RotatE
        see BaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -179,7 +179,7 @@ 
        besskge.scoring.RotatE
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -189,11 +189,11 @@ 
        besskge.scoring.RotatE
        p-norm reduction along embedding dimension.
        
 
        
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -204,7 +204,7 @@ 
        besskge.scoring.RotatE
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -215,16 +215,16 @@ 
        besskge.scoring.RotatE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -242,16 +242,16 @@ 
        besskge.scoring.RotatE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -269,16 +269,16 @@ 
        besskge.scoring.RotatE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -302,7 +302,7 @@ 
        besskge.scoring.RotatE
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.TranS.html b/generated/besskge.scoring.TranS.html
index b20e832..8fa3398 100644
--- a/generated/besskge.scoring.TranS.html
+++ b/generated/besskge.scoring.TranS.html
@@ -115,17 +115,17 @@ 
        besskge.scoring.TranS
 
        Parameters:
        
 
          
-
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
-
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
 
        • sharding (Sharding) – Entity sharding.
          • 
-
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
-
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
-
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
-
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
-
        • normalize_entities (bool) – If True, L2-normalize embeddings of head and tail entities as well as
+
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
+
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
+
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
+
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
+
        • normalize_entities (bool) – If True, L2-normalize embeddings of head and tail entities as well as
 tilde head and tail entities before multiplying. Default: True.
          • 
-
        • offset (float) – Offset applied to tilde entity embeddings. Default: 1.0.
          • 
-
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
+
        • offset (float) – Offset applied to tilde entity embeddings. Default: 1.0.
          • 
+
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
 
        
 
        
 
@@ -138,14 +138,14 @@ 
        besskge.scoring.TranS
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -157,7 +157,7 @@ 
        besskge.scoring.TranS
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -167,13 +167,13 @@ 
        besskge.scoring.TranS
        see BaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -181,7 +181,7 @@ 
        besskge.scoring.TranS
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -191,11 +191,11 @@ 
        besskge.scoring.TranS
        p-norm reduction along embedding dimension.
        
 
        
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -206,7 +206,7 @@ 
        besskge.scoring.TranS
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -217,16 +217,16 @@ 
        besskge.scoring.TranS
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -244,16 +244,16 @@ 
        besskge.scoring.TranS
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -271,16 +271,16 @@ 
        besskge.scoring.TranS
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -304,7 +304,7 @@ 
        besskge.scoring.TranS
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.TransE.html b/generated/besskge.scoring.TransE.html
index daf4998..f5a66a8 100644
--- a/generated/besskge.scoring.TransE.html
+++ b/generated/besskge.scoring.TransE.html
@@ -115,14 +115,14 @@ 
        besskge.scoring.TransE
 
        Parameters:
        
 
          
-
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
-
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
 
        • sharding (Sharding) – Entity sharding.
          • 
-
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
-
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
-
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
-
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
-
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
+
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
+
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
+
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
+
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
+
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
 
        
 
        
 
@@ -135,14 +135,14 @@ 
        besskge.scoring.TransE
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -154,7 +154,7 @@ 
        besskge.scoring.TransE
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -164,13 +164,13 @@ 
        besskge.scoring.TransE
        see BaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -178,7 +178,7 @@ 
        besskge.scoring.TransE
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -188,11 +188,11 @@ 
        besskge.scoring.TransE
        p-norm reduction along embedding dimension.
        
 
        
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -203,7 +203,7 @@ 
        besskge.scoring.TransE
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -214,16 +214,16 @@ 
        besskge.scoring.TransE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -241,16 +241,16 @@ 
        besskge.scoring.TransE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -268,16 +268,16 @@ 
        besskge.scoring.TransE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -301,7 +301,7 @@ 
        besskge.scoring.TransE
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.scoring.TripleRE.html b/generated/besskge.scoring.TripleRE.html
index 82f4b65..96308a8 100644
--- a/generated/besskge.scoring.TripleRE.html
+++ b/generated/besskge.scoring.TripleRE.html
@@ -115,18 +115,18 @@ 
        besskge.scoring.TripleRE
 
        Parameters:
        
 
          
-
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
-
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()
          • 
+
        • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()
          • 
 
        • sharding (Sharding) – Entity sharding.
          • 
-
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
-
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
-
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
-
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
-
        • normalize_entities (bool) – If True, L2-normalize head and tail entity embeddings before projecting.
+
        • n_relation_type (int) – Number of relation types in the knowledge graph.
          • 
+
        • embedding_size (int) – Size of entity and relation embeddings.
          • 
+
        • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.
          • 
+
        • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.
          • 
+
        • normalize_entities (bool) – If True, L2-normalize head and tail entity embeddings before projecting.
 Default: False.
          • 
-
        • u (float) – Offset factor for head/tail relation projections, as in TripleREv2.
+
        • u (float) – Offset factor for head/tail relation projections, as in TripleREv2.
 Default: 0.0 (no offset).
          • 
-
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
+
        • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False
          • 
 
        
 
        
 
@@ -139,14 +139,14 @@ 
        besskge.scoring.TripleRE
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (batch_size, embedding_size)
+
        • v1 (Tensor) – shape: (batch_size, embedding_size)
 Batch queries.
          • 
-
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
+
        • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size
 Negative embeddings to score against queries.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_neg) if
@@ -158,7 +158,7 @@ 
        besskge.scoring.TripleRE
 
        
-entity_embedding: Parameter
        
+entity_embedding: Parameter
 
        Entity embedding table
        
 
        
 
@@ -168,13 +168,13 @@ 
        besskge.scoring.TripleREBaseScoreFunction.score_triple()
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Parameters:
        
 
        
 
        
 
        
@@ -182,7 +182,7 @@ 
        besskge.scoring.TripleRE
 
        
-negative_sample_sharing: bool
        
+negative_sample_sharing: bool
 
        Share negative entities to construct negative samples
        
 
        
 
@@ -192,11 +192,11 @@ 
        besskge.scoring.TripleRE
 
        Parameters:
        
-
        v (Tensor) – shape: (*, embedding_size)
+
        v (Tensor) – shape: (*, embedding_size)
 The tensor to reduce.
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (*,)
@@ -207,7 +207,7 @@ 
        besskge.scoring.TripleRE
 
        
-relation_embedding: Parameter
        
+relation_embedding: Parameter
 
        Relation embedding table
        
 
        
 
@@ -218,16 +218,16 @@ 
        besskge.scoring.TripleRE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
+
        • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size
 Embeddings of head entities.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embedding of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_heads)
@@ -245,16 +245,16 @@ 
        besskge.scoring.TripleRE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
+
        • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size
 Embedding of tail entities.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size, B * n_tails)
@@ -272,16 +272,16 @@ 
        besskge.scoring.TripleRE
 
        Parameters:
        
 
          
-
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • head_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of head entities in batch.
          • 
-
        • relation_id (Tensor) – shape: (batch_size,)
+
        • relation_id (Tensor) – shape: (batch_size,)
 IDs of relation types in batch.
          • 
-
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
+
        • tail_emb (Tensor) – shape: (batch_size, embedding_size)
 Embeddings of tail entities in batch.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (batch_size,)
@@ -305,7 +305,7 @@ 
        besskge.scoring.TripleRE
        new_sharding (Sharding) – The new entity sharding.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
 
        
diff --git a/generated/besskge.sharding.PartitionedTripleSet.html b/generated/besskge.sharding.PartitionedTripleSet.html
index 098e606..ef39c15 100644
--- a/generated/besskge.sharding.PartitionedTripleSet.html
+++ b/generated/besskge.sharding.PartitionedTripleSet.html
@@ -114,16 +114,16 @@ 
        besskge.sharding.PartitionedTripleSetParameters:
 
        
 
        
 
@@ -135,11 +135,11 @@ 
        besskge.sharding.PartitionedTripleSetParameters:
 
          
 
        • dataset (KGDataset) – Knowledge graph dataset.
          • 
-
        • part (str) – The dataset part to shard.
          • 
+
        • part (str) – The dataset part to shard.
          • 
 
        • sharding (Sharding) – The entity sharding to use.
          • 
-
        • partition_mode (str) – The triple partition mode. Can be
+
        • partition_mode (str) – The triple partition mode. Can be
 “h_shard”, “t_shard”, “ht_shardpair”.
          • 
-
        • add_inverse_triples (bool) – 
          • 
+
        • add_inverse_triples (bool) – 
          • 
 
        
 
        
 
        Return type:
        
@@ -161,18 +161,18 @@ 
        besskge.sharding.PartitionedTripleSet
          
 
        • dataset (KGDataset) – Knowledge graph dataset.
          • 
 
        • sharding (Sharding) – The entity sharding to use.
          • 
-
        • queries (ndarray[Any, dtype[int32]]) – shape: (n_query, 2)
+
        • queries (ndarray[Any, dtype[int32]]) – shape: (n_query, 2)
 The set of (h, r) or (r, t) queries.
 Global IDs for entities/relations.
          • 
-
        • query_mode (str) – “hr” for (h,r,?) queries, “rt” for (?,r,t) queries.
          • 
-
        • ground_truth (Optional[ndarray[Any, dtype[int32]]]) – shape: (n_query,)
+
        • query_mode (str) – “hr” for (h,r,?) queries, “rt” for (?,r,t) queries.
          • 
+
        • ground_truth (Optional[ndarray[Any, dtype[int32]]]) – shape: (n_query,)
 If known, the global ID of the ground truth tail/head.
          • 
-
        • negative (Optional[ndarray[Any, dtype[int32]]]) – shape: (N, n_negative)
+
        • negative (Optional[ndarray[Any, dtype[int32]]]) – shape: (N, n_negative)
 Global IDs of negative entities to score against each query. This
 can be query-specific (N=n_query) or the same for all queries (N=1).
 Default: None (namely score each query against all entities in the
 graph).
          • 
-
        • negative_type (Optional[str]) – Score each query only against entities of a specific type. Default:
+
        • negative_type (Optional[str]) – Score each query only against entities of a specific type. Default:
 None (namely score each query against entities of any type).
          • 
 
        
 
        
@@ -187,7 +187,7 @@ 
        besskge.sharding.PartitionedTripleSet
 
        
-dummy: Optional[str]
        
+dummy: Optional[str]
 
        If set is constructed from (h,r,?) (resp. (?,r,t)) queries,
 dummy tails (resp. heads) are added to make pairs into triples.
 “head”, “tail”, “none”
        
@@ -195,28 +195,28 @@ 
        besskge.sharding.PartitionedTripleSet
 
        
-inverse_triples: bool
        
+inverse_triples: bool
 
        Whether the collection contains inverse triples (t,r_inv,h)
 for each regular triple (h,r,t)
        
 
        
 
 
        
 
        
-neg_heads: Optional[ndarray[Any, dtype[int32]]]
        
+neg_heads: Optional[ndarray[Any, dtype[int32]]]
 
        Global IDs of (possibly triple-specific) negative heads;
 int32[n_triple or 1, n_neg_heads]
        
 
        
 
 
        
 
        
-neg_tails: Optional[ndarray[Any, dtype[int32]]]
        
+neg_tails: Optional[ndarray[Any, dtype[int32]]]
 
        Global IDs of (possibly triple-specific) negative heads;
 int32[n_triple or 1, n_neg_tails]
        
 
        
 
 
        
 
        
-partition_mode: str
        
+partition_mode: str
 
        Partitioning criterion for triples;
 “h_shard”, “t_shard”, “ht_shardpair”
        
 
        
@@ -229,28 +229,28 @@ 
        besskge.sharding.PartitionedTripleSet
 
        
-triple_counts: ndarray[Any, dtype[int64]]
        
+triple_counts: ndarray[Any, dtype[int64]]
 
        Number of triples in each partition;
 int64[n_shard] or int64[n_shard, n_shard]
        
 
        
 
 
        
 
        
-triple_offsets: ndarray[Any, dtype[int64]]
        
+triple_offsets: ndarray[Any, dtype[int64]]
 
        Delimiting indices of ordered partitions;
 int64[n_shard] or int64[n_shard, n_shard]
        
 
        
 
 
        
 
        
-triple_sort_idx: ndarray[Any, dtype[int64]]
        
+triple_sort_idx: ndarray[Any, dtype[int64]]
 
        Sorting indices to order triples by partition;
 int64[n_triple]
        
 
        
 
 
        
 
        
-triples: ndarray[Any, dtype[int32]]
        
+triples: ndarray[Any, dtype[int32]]
 
        h/r/t IDs for triples ordered by partition.
 Local IDs for heads (resp. tails) and global IDs
 for tails (resp. heads) if partition_mode = “h_shard” (resp. “t_shard”);
@@ -260,7 +260,7 @@ 
        besskge.sharding.PartitionedTripleSet
 
        
-types: Optional[ndarray[Any, dtype[int32]]]
        
+types: Optional[ndarray[Any, dtype[int32]]]
 
        Entity type IDs of triple head/tail;
 int32[n_triple, {h_type, t_type}]
        
 
        
diff --git a/generated/besskge.sharding.Sharding.html b/generated/besskge.sharding.Sharding.html
index 071ff8c..a3793cc 100644
--- a/generated/besskge.sharding.Sharding.html
+++ b/generated/besskge.sharding.Sharding.html
@@ -104,13 +104,13 @@ 
        besskge.sharding.Sharding
 
        Parameters:
        
 
        
 
        
 
@@ -121,10 +121,10 @@ 
        besskge.sharding.Sharding
 
        Parameters:
        
 
          
-
        • n_entity (int) – Number of entities in the knowledge graph.
          • 
-
        • n_shard (int) – Number of shards.
          • 
-
        • seed (int) – Seed for random sharding.
          • 
-
        • type_offsets (Optional[ndarray[Any, dtype[int64]]]) – shape: (n_types,)
+
        • n_entity (int) – Number of entities in the knowledge graph.
          • 
+
        • n_shard (int) – Number of shards.
          • 
+
        • seed (int) – Seed for random sharding.
          • 
+
        • type_offsets (Optional[ndarray[Any, dtype[int64]]]) – shape: (n_types,)
 Global offsets of entity types. Default: None.
          • 
 
        
 
        
@@ -139,28 +139,28 @@ 
        besskge.sharding.Sharding
 
        
-entity_to_idx: ndarray[Any, dtype[int32]]
        
+entity_to_idx: ndarray[Any, dtype[int32]]
 
        Entity local ID on shard by global ID;
 int32[n_entity]
        
 
        
 
 
        
 
        
-entity_to_shard: ndarray[Any, dtype[int32]]
        
+entity_to_shard: ndarray[Any, dtype[int32]]
 
        Entity shard by global ID;
 int32[n_entity]
        
 
        
 
 
        
 
        
-entity_type_counts: Optional[ndarray[Any, dtype[int64]]]
        
+entity_type_counts: Optional[ndarray[Any, dtype[int64]]]
 
        Number of entities of each type on each shard;
 int64[n_shard, n_types]
        
 
        
 
 
        
 
        
-entity_type_offsets: Optional[ndarray[Any, dtype[int64]]]
        
+entity_type_offsets: Optional[ndarray[Any, dtype[int64]]]
 
        Offsets for entities of same type on each shared
 (entities remain clustered by type also locally);
 int64[n_shard, n_types]
        
@@ -172,7 +172,7 @@ 
        besskge.sharding.ShardingSharding object saved with Sharding.save().
        
 
        
 
        Parameters:
        
-
        path (Path) – Path to saved Sharding object.
        
+
        path (Path) – Path to saved Sharding object.
        
 
        
 
        Return type:
        
 
        Sharding
        
@@ -185,20 +185,20 @@ 
        besskge.sharding.Sharding
 
        
-property max_entity_per_shard: int
        
+property max_entity_per_shard: int
 
        Number of entities in a shard, after applying
 padding.
        
 
        
 
 
        
 
        
-property n_entity: int
        
+property n_entity: int
 
        Number of entities in the knowledge graph.
        
 
        
 
 
        
 
        
-n_shard: int
        
+n_shard: int
 
        Number of shards
        
 
        
 
@@ -208,24 +208,24 @@ 
        besskge.sharding.Sharding
 
        Parameters:
        
-
        out_file (Path) – Path to output file.
        
+
        out_file (Path) – Path to output file.
        
 
        
 
        Return type:
        
-
        None
        
+
        None
        
 
        
 
        
 
        
 
 
        
 
        
-shard_and_idx_to_entity: ndarray[Any, dtype[int32]]
        
+shard_and_idx_to_entity: ndarray[Any, dtype[int32]]
 
        Entity global ID by (shard, local_ID);
 int32[n_shard, max_entity_per_shard]
        
 
        
 
 
        
 
        
-shard_counts: ndarray[Any, dtype[int64]]
        
+shard_counts: ndarray[Any, dtype[int64]]
 
        Number of true entities (excluding padding) in each shard;
 int64[n_shard]
        
 
        
diff --git a/generated/besskge.utils.complex_multiplication.html b/generated/besskge.utils.complex_multiplication.html
index 5ff7c79..e3a0e96 100644
--- a/generated/besskge.utils.complex_multiplication.html
+++ b/generated/besskge.utils.complex_multiplication.html
@@ -105,14 +105,14 @@ 
        besskge.utils.complex_multiplication
 
        Parameters:
        
 
          
-
        • v1 (Tensor) – shape: (a, 2*e)
+
        • v1 (Tensor) – shape: (a, 2*e)
 v1[:,:e] real part, v1[:,e:] imaginary part.
          • 
-
        • v2 (Tensor) – shape: (a, 2*e)
+
        • v2 (Tensor) – shape: (a, 2*e)
 v2[:,:e] real part, v2[:,e:] imaginary part.
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (a, 2*e)
diff --git a/generated/besskge.utils.complex_rotation.html b/generated/besskge.utils.complex_rotation.html
index b339124..16454d7 100644
--- a/generated/besskge.utils.complex_rotation.html
+++ b/generated/besskge.utils.complex_rotation.html
@@ -106,15 +106,15 @@ 
        besskge.utils.complex_rotation
 
        Parameters:
        
 
          
-
        • v (Tensor) – shape: (a, 2*e)
+
        • v (Tensor) – shape: (a, 2*e)
 Complex tensor to rotate:
 v[:,:e] real part, v[:,e:] imaginary part.
          • 
-
        • r (Tensor) – shape: (a, e)
+
        • r (Tensor) – shape: (a, e)
 Rotate v[k] by \(e^{i \pi r[k]}\)
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (a, 2*e)
diff --git a/generated/besskge.utils.gather_indices.html b/generated/besskge.utils.gather_indices.html
index e70264f..b556058 100644
--- a/generated/besskge.utils.gather_indices.html
+++ b/generated/besskge.utils.gather_indices.html
@@ -101,17 +101,17 @@ 
        besskge.utils.gather_indices
 
        
 besskge.utils.gather_indices(x, index)[source]
        
-
        IPU-friendly gather function like torch.take_along_dim()
+
        IPU-friendly gather function like torch.take_along_dim()
 for 2-dimensional tensors (indices along dim=1).
        
 
        
 
        Parameters:
        
 
          
-
        • x (Tensor) – shape: (a, e)
          • 
-
        • index (Tensor) – shape: (b, k)
          • 
+
        • x (Tensor) – shape: (a, e)
          • 
+
        • index (Tensor) – shape: (b, k)
          • 
 
        
 
        
 
        Return type:
        
-
        Tensor
        
+
        Tensor
        
 
        
 
        Returns:
        
 
        shape: (b, k)
diff --git a/generated/besskge.utils.html b/generated/besskge.utils.html
index 6a9d15f..4e7fb02 100644
--- a/generated/besskge.utils.html
+++ b/generated/besskge.utils.html
@@ -105,7 +105,7 @@
 
        Batched rotation by unitary tensors.
        
 
 
        gather_indices(x, index)
        
-
        IPU-friendly gather function like torch.take_along_dim() for 2-dimensional tensors (indices along dim=1).
        
+
        IPU-friendly gather function like torch.take_along_dim() for 2-dimensional tensors (indices along dim=1).