pykeen · cthoyt · Apr 20, 2025 · Apr 20, 2025 · Apr 20, 2025 · Apr 21, 2025
diff --git a/pyproject.toml b/pyproject.toml
@@ -176,6 +176,7 @@ skip = [
 # MyPy, see https://mypy.readthedocs.io/en/stable/config_file.html
 [tool.mypy]
 plugins = [
+    "numpy.typing.mypy_plugin",
 ]
 
 # Doc8, see https://doc8.readthedocs.io/en/stable/readme.html#ini-file-usage

diff --git a/src/pykeen/checkpoints/keeper.py b/src/pykeen/checkpoints/keeper.py
@@ -120,5 +120,6 @@ def __call__(self, steps: Sequence[int]) -> Iterator[int]:
 
 #: a resolver for checkpoint keepers
 keeper_resolver: ClassResolver[CheckpointKeeper] = ClassResolver.from_subclasses(
-    CheckpointKeeper, default=CheckpointKeeper
+    CheckpointKeeper,  # type:ignore[type-abstract]
+    default=CheckpointKeeper,
 )
diff --git a/src/pykeen/checkpoints/schedule.py b/src/pykeen/checkpoints/schedule.py
@@ -89,4 +89,7 @@ def __call__(self, step: int) -> bool:
 
 
 #: a resolver for checkpoint schedules
-schedule_resolver = ClassResolver.from_subclasses(base=CheckpointSchedule, default=EveryCheckpointSchedule)
+schedule_resolver = ClassResolver.from_subclasses(
+    base=CheckpointSchedule,  # type:ignore[type-abstract]
+    default=EveryCheckpointSchedule,
+)
diff --git a/src/pykeen/contrib/lightning.py b/src/pykeen/contrib/lightning.py
@@ -250,7 +250,7 @@ def _dataloader(self, triples_factory: CoreTriplesFactory, shuffle: bool = False
 
 #: A resolver for PyTorch Lightning training modules
 lit_module_resolver: ClassResolver[LitModule] = ClassResolver.from_subclasses(
-    base=LitModule,
+    base=LitModule,  # type:ignore[type-abstract]
     default=SLCWALitModule,
     # note: since this file is executed via __main__, its module name is replaced by __name__
     #       hence, the two classes' fully qualified names start with "_" and are considered private

diff --git a/src/pykeen/datasets/ea/combination.py b/src/pykeen/datasets/ea/combination.py
@@ -36,6 +36,7 @@
     "CollapseGraphPairCombinator",
     # Data Structures
     "ProcessedTuple",
+    "graph_combinator_resolver",
 ]
 
 logger = logging.getLogger(__name__)
@@ -491,7 +492,8 @@ def process(
         )
 
 
+#: A resolver for graph combinatiors
 graph_combinator_resolver: ClassResolver[GraphPairCombinator] = ClassResolver.from_subclasses(
-    base=GraphPairCombinator,
+    base=GraphPairCombinator,  # type:ignore[type-abstract]
     default=ExtraRelationGraphPairCombinator,
 )
diff --git a/src/pykeen/evaluation/__init__.py b/src/pykeen/evaluation/__init__.py
@@ -30,9 +30,11 @@
 
 #: A resolver for evaluators
 evaluator_resolver: ClassResolver[Evaluator] = ClassResolver.from_subclasses(
-    base=Evaluator,
+    base=Evaluator,  # type:ignore[type-abstract]
     default=RankBasedEvaluator,
 )
 
 #: A resolver for metric results
-metric_resolver: ClassResolver[MetricResults] = ClassResolver.from_subclasses(MetricResults)
+metric_resolver: ClassResolver[MetricResults] = ClassResolver.from_subclasses(
+    MetricResults,  # type:ignore[type-abstract]
+)
diff --git a/src/pykeen/experiments/validate.py b/src/pykeen/experiments/validate.py
@@ -46,7 +46,7 @@
     type[Model],
     Optional[type[Model]],  # noqa:UP007
     Union[str, Callable[[FloatTensor], FloatTensor]],  # noqa:UP007
-    Hint[nn.Module],
+    Hint[nn.Module],  # type:ignore
 }
 _SKIP_EXTRANEOUS = {
     "predict_with_sigmoid",

diff --git a/src/pykeen/inverse.py b/src/pykeen/inverse.py
@@ -71,6 +71,6 @@ def is_inverse(self, ids: LongTensor) -> BoolTensor:  # noqa: D102
 
 #: A resolver for relation inverter protocols
 relation_inverter_resolver: Resolver[RelationInverter] = Resolver.from_subclasses(
-    RelationInverter,
+    RelationInverter,  # type:ignore[type-abstract]
     default=DefaultRelationInverter,
 )
diff --git a/src/pykeen/losses.py b/src/pykeen/losses.py
@@ -1748,15 +1748,15 @@ def forward(self, x: FloatTensor, target: FloatTensor, weight: FloatTensor | Non
 
 #: A resolver for loss modules
 loss_resolver: ClassResolver[Loss] = ClassResolver.from_subclasses(
-    Loss,
+    Loss,  # type:ignore[type-abstract]
     default=MarginRankingLoss,
     skip={
-        PairwiseLoss,
-        PointwiseLoss,
-        SetwiseLoss,
+        PairwiseLoss,  # type:ignore[type-abstract]
+        PointwiseLoss,  # type:ignore[type-abstract]
+        SetwiseLoss,  # type:ignore[type-abstract]
         DeltaPointwiseLoss,
         MarginPairwiseLoss,
-        AdversarialLoss,
+        AdversarialLoss,  # type:ignore[type-abstract]
     },
 )
 for _name, _cls in loss_resolver.lookup_dict.items():

diff --git a/src/pykeen/metrics/classification.py b/src/pykeen/metrics/classification.py
@@ -756,7 +756,10 @@ def extract_from_confusion_matrix(self, matrix: numpy.ndarray) -> float:  # noqa
 
 #: A resolver for classification metrics
 classification_metric_resolver: ClassResolver[ClassificationMetric] = ClassResolver.from_subclasses(
-    base=ClassificationMetric,
+    base=ClassificationMetric,  # type:ignore[type-abstract]
     default=AveragePrecisionScore,
-    skip={BinarizedClassificationMetric, ConfusionMatrixClassificationMetric},
+    skip={
+        BinarizedClassificationMetric,  # type:ignore[type-abstract]
+        ConfusionMatrixClassificationMetric,  # type:ignore[type-abstract]
+    },
 )
diff --git a/src/pykeen/metrics/ranking.py b/src/pykeen/metrics/ranking.py
@@ -1667,9 +1667,14 @@ class AdjustedGeometricMeanRankIndex(ReindexedMetric):
 
 
 rank_based_metric_resolver: ClassResolver[RankBasedMetric] = ClassResolver.from_subclasses(
-    base=RankBasedMetric,
+    base=RankBasedMetric,  # type:ignore[type-abstract]
     default=InverseHarmonicMeanRank,  # mrr
-    skip={ExpectationNormalizedMetric, ReindexedMetric, ZMetric, DerivedRankBasedMetric},
+    skip={
+        ExpectationNormalizedMetric,  # type:ignore[type-abstract]
+        ReindexedMetric,  # type:ignore[type-abstract]
+        ZMetric,  # type:ignore[type-abstract]
+        DerivedRankBasedMetric,  # type:ignore[type-abstract]
+    },
 )
 """The rank-based metric resolver allows for the lookup and instantiation of classes
 deriving from :class:`RankBasedMetric` via the :mod:`class_resolver`.

diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py
@@ -152,16 +152,16 @@
 
 #: A resolver for knowledge graph embedding models
 model_resolver: ClassResolver[Model] = ClassResolver.from_subclasses(
-    base=Model,
+    base=Model,  # type:ignore[type-abstract]
     skip={
         # Abstract Models
-        _NewAbstractModel,
+        _NewAbstractModel,  # type:ignore[type-abstract]
         # We might be able to relax this later
         ERModel,
         InductiveERModel,
         LiteralModel,
         # baseline models behave differently
-        EvaluationOnlyModel,
-        *get_subclasses(EvaluationOnlyModel),
+        EvaluationOnlyModel,  # type:ignore[type-abstract]
+        *get_subclasses(EvaluationOnlyModel),  # type:ignore[type-abstract]
     },
 )
diff --git a/src/pykeen/nn/message_passing.py b/src/pykeen/nn/message_passing.py
@@ -507,7 +507,8 @@ def forward(
 
 
 decomposition_resolver: ClassResolver[Decomposition] = ClassResolver.from_subclasses(
-    base=Decomposition, default=BasesDecomposition
+    base=Decomposition,  # type:ignore[type-abstract]
+    default=BasesDecomposition,
 )
 
 
@@ -545,6 +546,10 @@ class RGCNRepresentation(Representation):
         github: https://github.com/MichSchli/RelationPrediction
     """
 
+    sources: LongTensor
+    targets: LongTensor
+    edge_types: LongTensor
+
     @update_docstring_with_resolver_keys(
         ResolverKey("entity_representations", resolver="pykeen.nn.representation_resolver"),
         ResolverKey("activation", resolver="class_resolver.contrib.torch.activation_resolver"),

diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py
@@ -3440,8 +3440,8 @@ def __init__(
         num_heads: int = 8,
         dropout: float = 0.1,
         dim_feedforward: int = 2048,
-        position_initializer: HintOrType[Initializer] = xavier_normal_,
-    ):
+        position_initializer: HintOrType[Initializer] = None,
+    ) -> None:
         """
         Initialize the module.
 
@@ -3470,6 +3470,8 @@ def __init__(
             ),
             num_layers=num_layers,
         )
+        if position_initializer is None:
+            position_initializer = xavier_normal_
         self.position_embeddings = nn.Parameter(position_initializer(torch.empty(2, input_dim)))
         self.final = nn.Linear(input_dim, input_dim, bias=True)
 
@@ -3875,9 +3877,9 @@ def forward(self, h: FloatTensor, r: tuple[FloatTensor, FloatTensor, FloatTensor
 
 #: A resolver for stateful interaction functions
 interaction_resolver: ClassResolver[Interaction] = ClassResolver.from_subclasses(
-    Interaction,
+    Interaction,  # type:ignore[type-abstract]
     skip={
-        NormBasedInteraction,
+        NormBasedInteraction,  # type:ignore[type-abstract]
         MonotonicAffineTransformationInteraction,
         ClampedInteraction,
         DirectionAverageInteraction,

diff --git a/src/pykeen/nn/node_piece/anchor_selection.py b/src/pykeen/nn/node_piece/anchor_selection.py
@@ -8,6 +8,7 @@
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence
+from typing import Any
 
 import numpy
 import torch
@@ -149,7 +150,7 @@ class PageRankAnchorSelection(SingleSelection):
     def __init__(
         self,
         num_anchors: int = 32,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """Initialize the selection strategy.
 
@@ -201,7 +202,7 @@ def __init__(
         selections: Sequence[HintOrType[AnchorSelection]],
         ratios: None | float | Sequence[float] = None,
         selections_kwargs: OneOrSequence[OptionalKwargs] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> None:
         """Initialize the selection strategy.
 
@@ -257,7 +258,7 @@ def __call__(
 
 #: A resolver for NodePiece anchor selectors
 anchor_selection_resolver: ClassResolver[AnchorSelection] = ClassResolver.from_subclasses(
-    base=AnchorSelection,
+    base=AnchorSelection,  # type:ignore[type-abstract]
     default=DegreeAnchorSelection,
     skip={SingleSelection},
 )
diff --git a/src/pykeen/nn/node_piece/loader.py b/src/pykeen/nn/node_piece/loader.py
@@ -93,6 +93,6 @@ def __call__(self, path: pathlib.Path) -> tuple[Mapping[int, Collection[int]], i
 
 #: A resolver for NodePiece precomputed tokenizer loaders
 precomputed_tokenizer_loader_resolver: ClassResolver[PrecomputedTokenizerLoader] = ClassResolver.from_subclasses(
-    base=PrecomputedTokenizerLoader,
+    base=PrecomputedTokenizerLoader,  # type:ignore[type-abstract]
     default=GalkinPrecomputedTokenizerLoader,
 )
diff --git a/src/pykeen/nn/node_piece/representation.py b/src/pykeen/nn/node_piece/representation.py
@@ -2,7 +2,7 @@
 
 import logging
 import pathlib
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Sequence
 from typing import NamedTuple
 
 import torch
@@ -257,6 +257,8 @@ class NodePieceRepresentation(CombinedRepresentation):
         github: https://github.com/migalkin/NodePiece
     """
 
+    base: Sequence[TokenizationRepresentation]
+
     @update_docstring_with_resolver_keys(
         ResolverKey("token_representations", resolver="pykeen.nn.representation_resolver"),
         ResolverKey("tokenizers", resolver="pykeen.nn.node_piece.tokenizer_resolver"),
@@ -275,7 +277,7 @@ def __init__(
         aggregation_kwargs: OptionalKwargs = None,
         max_id: int | None = None,
         **kwargs,
-    ):
+    ) -> None:
         """
         Initialize the representation.
 
@@ -322,12 +324,12 @@ def __init__(
             # inverse triples are created afterwards implicitly
             mapped_triples = mapped_triples[mapped_triples[:, 1] < triples_factory.real_num_relations]
 
-        token_representations, token_representations_kwargs, num_tokens = broadcast_upgrade_to_sequences(
+        token_representations_, token_representations_kwargs_, num_tokens = broadcast_upgrade_to_sequences(
             token_representations, token_representations_kwargs, num_tokens
         )
 
         # tokenize
-        token_representations = [
+        base = [
             TokenizationRepresentation.from_tokenizer(
                 tokenizer=tokenizer_inst,
                 num_tokens=num_tokens_,
@@ -339,8 +341,8 @@ def __init__(
             )
             for tokenizer_inst, token_representation, token_representation_kwargs, num_tokens_ in zip(
                 tokenizer_resolver.make_many(queries=tokenizers, kwargs=tokenizers_kwargs),
-                token_representations,
-                token_representations_kwargs,
+                token_representations_,
+                token_representations_kwargs_,
                 num_tokens,
                 strict=False,
             )
@@ -349,18 +351,18 @@ def __init__(
         # Create an MLP for string aggregation
         if aggregation == "mlp":
             # note: the token representations' shape includes the number of tokens as leading dim
-            embedding_dim = token_representations[0].shape[1]
+            embedding_dim = base[0].shape[1]
             aggregation = ConcatMLP(
                 input_dim=embedding_dim * sum(num_tokens),
                 output_dim=embedding_dim,
             )
 
         super().__init__(
             max_id=triples_factory.num_entities,
-            base=token_representations,
+            base=base,
             combination=ConcatAggregationCombination,
             combination_kwargs=dict(
-                aggregation=aggregation, aggregation_kwargs=aggregation_kwargs, dim=-len(token_representations[0].shape)
+                aggregation=aggregation, aggregation_kwargs=aggregation_kwargs, dim=-len(base[0].shape)
             ),
             **kwargs,
         )

diff --git a/src/pykeen/nn/node_piece/tokenization.py b/src/pykeen/nn/node_piece/tokenization.py
@@ -5,6 +5,7 @@
 from abc import abstractmethod
 from collections import defaultdict
 from collections.abc import Collection, Mapping
+from typing import Any
 
 import more_itertools
 import numpy
@@ -126,15 +127,15 @@ def _call(
         num_tokens: int,
         num_entities: int,
     ) -> tuple[int, LongTensor]:
-        edge_index = edge_index.numpy()
+        edge_index_np = edge_index.numpy()
         # select anchors
         logger.info(f"Selecting anchors according to {self.anchor_selection}")
-        anchors = self.anchor_selection(edge_index=edge_index)
+        anchors = self.anchor_selection(edge_index=edge_index_np)
         if len(numpy.unique(anchors)) < len(anchors):
             logger.warning(f"Only {len(numpy.unique(anchors))} out of {len(anchors)} anchors are unique")
         # find closest anchors
         logger.info(f"Searching closest anchors with {self.searcher}")
-        tokens = self.searcher(edge_index=edge_index, anchors=anchors, k=num_tokens, num_entities=num_entities)
+        tokens = self.searcher(edge_index=edge_index_np, anchors=anchors, k=num_tokens, num_entities=num_entities)
         num_empty = (tokens < 0).all(axis=1).sum()
         if num_empty > 0:
             logger.warning(
@@ -165,7 +166,7 @@ class MetisAnchorTokenizer(AnchorTokenizer):
     http://glaros.dtc.umn.edu/gkhome/metis/metis/overview
     """
 
-    def __init__(self, num_partitions: int = 2, device: DeviceHint = None, **kwargs):
+    def __init__(self, num_partitions: int = 2, device: DeviceHint = None, **kwargs: Any) -> None:
         """Initialize the tokenizer.
 
         :param num_partitions: the number of partitions obtained through Metis.
@@ -280,7 +281,7 @@ def __init__(
         pool: Mapping[int, Collection[int]] | None = None,
         randomize_selection: bool = False,
         loader: HintOrType[PrecomputedTokenizerLoader] = None,
-    ):
+    ) -> None:
         r"""Initialize the tokenizer.
 
         .. note::
@@ -330,6 +331,6 @@ def __call__(
 
 #: A resolver for NodePiece tokenizers
 tokenizer_resolver: ClassResolver[Tokenizer] = ClassResolver.from_subclasses(
-    base=Tokenizer,
+    base=Tokenizer,  # type:ignore[type-abstract]
     default=RelationTokenizer,
 )
diff --git a/src/pykeen/nn/text/encoder.py b/src/pykeen/nn/text/encoder.py
@@ -235,6 +235,6 @@ def forward_normalized(self, texts: Sequence[str]) -> FloatTensor:  # noqa: D102
 #: for :class:`CharacterEmbeddingTextEncoder` or 'transformer' for
 #: :class:`TransformerTextEncoder`.
 text_encoder_resolver: ClassResolver[TextEncoder] = ClassResolver.from_subclasses(
-    base=TextEncoder,
+    base=TextEncoder,  # type:ignore[type-abstract]
     default=CharacterEmbeddingTextEncoder,
 )
diff --git a/src/pykeen/sampling/__init__.py b/src/pykeen/sampling/__init__.py
@@ -130,6 +130,6 @@
 
 #: A resolver for negative samplers
 negative_sampler_resolver: ClassResolver[NegativeSampler] = ClassResolver.from_subclasses(
-    NegativeSampler,
+    NegativeSampler,  # type:ignore[type-abstract]
     default=BasicNegativeSampler,
 )