⚖️🗺️ Fix estimate diversity function (pykeen#1242)

Closes pykeen#1241 Fixed a bug in NodePieceRepresentation's estimate_diversity function --------- Co-authored-by: Max Berrendorf <max.berrendorf@gmail.com> Co-authored-by: Charles Tapley Hoyt <cthoyt@gmail.com>
nicolafan · Mar 23, 2023 · b3ef9d1 · b3ef9d1
1 parent b253678
commit b3ef9d1
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/src/pykeen/nn/node_piece/representation.py b/src/pykeen/nn/node_piece/representation.py
@@ -348,13 +348,11 @@ def estimate_diversity(self) -> HashDiversityInfo:
         .. seealso:: https://github.com/pykeen/pykeen/pull/896
         """
         # unique hashes per representation
-        uniques_per_representation = [
-            tokens.assignment.unique(dim=0).shape[0] / self.max_id for tokens in self.token_representations
-        ]
+        uniques_per_representation = [tokens.assignment.unique(dim=0).shape[0] / self.max_id for tokens in self.base]
 
         # unique hashes if we concatenate all representations together
         unnormalized_uniques_total = torch.unique(
-            torch.cat([tokens.assignment for tokens in self.token_representations], dim=-1), dim=0
+            torch.cat([tokens.assignment for tokens in self.base], dim=-1), dim=0
         ).shape[0]
         return HashDiversityInfo(
             uniques_per_representation=uniques_per_representation,

diff --git a/tests/cases.py b/tests/cases.py
@@ -2069,6 +2069,12 @@ def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMa
         kwargs.pop("max_id")
         return kwargs
 
+    def test_estimate_diversity(self):
+        """Test estimating diversity."""
+        diversity = self.instance.estimate_diversity()
+        assert len(diversity.uniques_per_representation) == len(self.instance.base)
+        assert 0.0 <= diversity.uniques_total <= 1.0
+
 
 class EvaluationLoopTestCase(GenericTestCase[pykeen.evaluation.evaluation_loop.EvaluationLoop]):
     """Tests for evaluation loops."""