Jaxtyping Shape Contracts

CLAUDE.md

@@ -140,6 +140,22 @@ development.
 - Use `Final` for constants. Use `@dataclass(frozen=True)` for immutable data containers when named fields and a stable
   shape add real clarity; do not introduce a dataclass for tiny internal-only plumbing.
 - Always annotate empty containers: `names: list[str] = []` not `names = []`.
+- A Shape Contract is a runtime-checkable Jaxtyping dtype and shape annotation at a stable API boundary.
+- Use Jaxtyping annotations for tensors crossing loader or sampler boundaries, public model `forward` or `decode`
+  methods, and loss interfaces. Do not add them to internal tensor operations, dynamic PyG or TorchRec keyed containers,
+  or low-level message-passing operations unless they clarify a stable boundary.
+- Use an exact dtype such as `Int64`, `Int32`, `Float32`, or `UInt8` only when the boundary guarantees it. Keep `Float`
+  for model and loss boundaries that intentionally support mixed precision.
+- Reuse axis names when dimensions must match across annotations. Use `_name` when a dimension must not bind to another
+  annotation, `_` when its meaning is unknown, and `#name` when size `1` is valid because the dimension supports PyTorch
+  broadcasting. Use numeric dimensions only when the size is guaranteed. Use `...` or `*name` only when variable rank is
+  part of the boundary contract. Use `{expression}` for an exact dimension derived from a runtime argument or instance
+  configuration only when it adds a useful boundary contract. Whitespace separates axes; do not add leading or trailing
+  whitespace.
+- Unit, integration, and end-to-end test launchers install runtime Shape Contract checking before test discovery.
+  Typeguard checks every shape-bearing tensor value in annotated containers in `gigl` and `examples` when a contracted
+  call executes. Arguments are checked before execution and returns afterwards; an uncaught `jaxtyping.TypeCheckError`
+  fails the test command.
 
 ### Docstrings
 

Makefile

@@ -248,79 +248,79 @@ push_dev_workbench_docker_image: compile_jars
 run_cora_nalp_e2e_test: compiled_pipeline_path:=${GIGL_E2E_TEST_COMPILED_PIPELINE_PATH}
 run_cora_nalp_e2e_test: compile_gigl_kubeflow_pipeline
 run_cora_nalp_e2e_test:
-	uv run python tests/e2e_tests/e2e_test.py \
+	uv run python -m tests.e2e_tests.e2e_test \
 		--compiled_pipeline_path=$(compiled_pipeline_path) \
 		--test_spec_uri="tests/e2e_tests/e2e_tests.yaml" \
 		--test_names="cora_nalp_test"
 
 run_cora_snc_e2e_test: compiled_pipeline_path:=${GIGL_E2E_TEST_COMPILED_PIPELINE_PATH}
 run_cora_snc_e2e_test: compile_gigl_kubeflow_pipeline
 run_cora_snc_e2e_test:
-	uv run python tests/e2e_tests/e2e_test.py \
+	uv run python -m tests.e2e_tests.e2e_test \
 		--compiled_pipeline_path=$(compiled_pipeline_path) \
 		--test_spec_uri="tests/e2e_tests/e2e_tests.yaml" \
 		--test_names="cora_snc_test"
 
 run_cora_udl_e2e_test: compiled_pipeline_path:=${GIGL_E2E_TEST_COMPILED_PIPELINE_PATH}
 run_cora_udl_e2e_test: compile_gigl_kubeflow_pipeline
 run_cora_udl_e2e_test:
-	uv run python tests/e2e_tests/e2e_test.py \
+	uv run python -m tests.e2e_tests.e2e_test \
 		--compiled_pipeline_path=$(compiled_pipeline_path) \
 		--test_spec_uri="tests/e2e_tests/e2e_tests.yaml" \
 		--test_names="cora_udl_test"
 
 run_dblp_nalp_e2e_test: compiled_pipeline_path:=${GIGL_E2E_TEST_COMPILED_PIPELINE_PATH}
 run_dblp_nalp_e2e_test: compile_gigl_kubeflow_pipeline
 run_dblp_nalp_e2e_test:
-	uv run python tests/e2e_tests/e2e_test.py \
+	uv run python -m tests.e2e_tests.e2e_test \
 		--compiled_pipeline_path=$(compiled_pipeline_path) \
 		--test_spec_uri="tests/e2e_tests/e2e_tests.yaml" \
 		--test_names="dblp_nalp_test"
 
 run_hom_cora_sup_e2e_test: compiled_pipeline_path:=${GIGL_E2E_TEST_COMPILED_PIPELINE_PATH}
 run_hom_cora_sup_e2e_test: compile_gigl_kubeflow_pipeline
 run_hom_cora_sup_e2e_test:
-	uv run python tests/e2e_tests/e2e_test.py \
+	uv run python -m tests.e2e_tests.e2e_test \
 		--compiled_pipeline_path=$(compiled_pipeline_path) \
 		--test_spec_uri="tests/e2e_tests/e2e_tests.yaml" \
 		--test_names="hom_cora_sup_test"
 
 run_het_dblp_sup_e2e_test: compiled_pipeline_path:=${GIGL_E2E_TEST_COMPILED_PIPELINE_PATH}
 run_het_dblp_sup_e2e_test: compile_gigl_kubeflow_pipeline
 run_het_dblp_sup_e2e_test:
-	uv run python tests/e2e_tests/e2e_test.py \
+	uv run python -m tests.e2e_tests.e2e_test \
 		--compiled_pipeline_path=$(compiled_pipeline_path) \
 		--test_spec_uri="tests/e2e_tests/e2e_tests.yaml" \
 		--test_names="het_dblp_sup_test"
 
 run_hom_cora_sup_gs_e2e_test: compiled_pipeline_path:=${GIGL_E2E_TEST_COMPILED_PIPELINE_PATH}
 run_hom_cora_sup_gs_e2e_test: compile_gigl_kubeflow_pipeline
 run_hom_cora_sup_gs_e2e_test:
-	uv run python tests/e2e_tests/e2e_test.py \
+	uv run python -m tests.e2e_tests.e2e_test \
 		--compiled_pipeline_path=$(compiled_pipeline_path) \
 		--test_spec_uri="tests/e2e_tests/e2e_tests.yaml" \
 		--test_names="hom_cora_sup_gs_test"
 
 run_het_dblp_sup_gs_e2e_test: compiled_pipeline_path:=${GIGL_E2E_TEST_COMPILED_PIPELINE_PATH}
 run_het_dblp_sup_gs_e2e_test: compile_gigl_kubeflow_pipeline
 run_het_dblp_sup_gs_e2e_test:
-	uv run python tests/e2e_tests/e2e_test.py \
+	uv run python -m tests.e2e_tests.e2e_test \
 		--compiled_pipeline_path=$(compiled_pipeline_path) \
 		--test_spec_uri="tests/e2e_tests/e2e_tests.yaml" \
 		--test_names="het_dblp_sup_gs_test"
 
 run_hom_cora_snc_e2e_test: compiled_pipeline_path:=${GIGL_E2E_TEST_COMPILED_PIPELINE_PATH}
 run_hom_cora_snc_e2e_test: compile_gigl_kubeflow_pipeline
 run_hom_cora_snc_e2e_test:
-	uv run python tests/e2e_tests/e2e_test.py \
+	uv run python -m tests.e2e_tests.e2e_test \
 		--compiled_pipeline_path=$(compiled_pipeline_path) \
 		--test_spec_uri="tests/e2e_tests/e2e_tests.yaml" \
 		--test_names="hom_cora_snc_test"
 
 run_all_e2e_tests: compiled_pipeline_path:=${GIGL_E2E_TEST_COMPILED_PIPELINE_PATH}
 run_all_e2e_tests: compile_gigl_kubeflow_pipeline
 run_all_e2e_tests:
-	uv run python tests/e2e_tests/e2e_test.py \
+	uv run python -m tests.e2e_tests.e2e_test \
 		--compiled_pipeline_path=$(compiled_pipeline_path) \
 		--test_spec_uri="tests/e2e_tests/e2e_tests.yaml"
 

examples/link_prediction/graph_store/heterogeneous_training.py

@@ -86,6 +86,7 @@
 import torch
 import torch.distributed
 import torch.multiprocessing as mp
+from jaxtyping import Float
 from torch_geometric.data import HeteroData
 
 from examples.link_prediction.models import init_example_gigl_heterogeneous_model
@@ -263,7 +264,7 @@ def _compute_loss(
     supervision_edge_type: EdgeType,
     edge_dir: str,
     device: torch.device,
-) -> torch.Tensor:
+) -> Float[torch.Tensor, ""]:
     """
     With the provided model and loss function, computes the forward pass on the main batch data and random negative data.
     Args:

examples/link_prediction/graph_store/homogeneous_training.py

@@ -130,6 +130,7 @@
 import torch
 import torch.distributed
 import torch.multiprocessing as mp
+from jaxtyping import Float
 from torch_geometric.data import Data
 
 from examples.link_prediction.models import init_example_gigl_homogeneous_model
@@ -287,7 +288,7 @@ def _compute_loss(
     random_negative_data: Data,
     loss_fn: RetrievalLoss,
     device: torch.device,
-) -> torch.Tensor:
+) -> Float[torch.Tensor, ""]:
     """
     With the provided model and loss function, computes the forward pass on the main batch data and random negative data.
     Args:

examples/link_prediction/heterogeneous_training.py

@@ -37,6 +37,7 @@
 import torch
 import torch.distributed
 import torch.multiprocessing as mp
+from jaxtyping import Float
 from torch_geometric.data import HeteroData
 
 import gigl.distributed.utils
@@ -190,7 +191,7 @@ def _compute_loss(
     loss_fn: RetrievalLoss,
     supervision_edge_type: EdgeType,
     device: torch.device,
-) -> torch.Tensor:
+) -> Float[torch.Tensor, ""]:
     """
     With the provided model and loss function, computes the forward pass on the main batch data and random negative data.
     Args:

examples/link_prediction/homogeneous_training.py

@@ -32,6 +32,7 @@
 import torch
 import torch.distributed
 import torch.multiprocessing as mp
+from jaxtyping import Float
 from torch_geometric.data import Data
 
 import gigl.distributed.utils
@@ -177,7 +178,7 @@ def _compute_loss(
     random_negative_data: Data,
     loss_fn: RetrievalLoss,
     device: torch.device,
-) -> torch.Tensor:
+) -> Float[torch.Tensor, ""]:
     """
     With the provided model and loss function, computes the forward pass on the main batch data and random negative data.
     Args:

examples/node_classification/models.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.nn as nn
+from jaxtyping import Float
 from torch.nn.parallel import DistributedDataParallel
 from torch_geometric.data import Data
 from typing_extensions import Self
@@ -36,7 +37,9 @@ def encoder(self) -> nn.Module:
     def head(self) -> nn.Module:
         return self._head
 
-    def forward(self, data: Data, device: torch.device) -> torch.Tensor:
+    def forward(
+        self, data: Data, device: torch.device
+    ) -> Float[torch.Tensor, "nodes classes"]:
         """
         Runs the encoder then the classifier head on a sampled subgraph batch.
 

examples/tutorial/KDD_2025/heterogeneous_training.py

@@ -41,6 +41,7 @@
 from typing import Literal
 
 import torch
+from jaxtyping import Float
 from torch.nn.parallel import DistributedDataParallel
 from torch_geometric.data import HeteroData
 
@@ -70,7 +71,7 @@
 FANOUT = [10, 10]
 
 
-def compute_loss(model: torch.nn.Module, data: HeteroData) -> torch.Tensor:
+def compute_loss(model: torch.nn.Module, data: HeteroData) -> Float[torch.Tensor, ""]:
     main_out: dict[str, torch.Tensor] = model(data.x_dict, data.edge_index_dict)
     anchor_nodes = torch.arange(data[QUERY_NODE_TYPE].batch_size).repeat_interleave(
         torch.tensor([len(v) for v in data.y_positive.values()])

gigl/common/data/dataloaders.py

@@ -7,6 +7,7 @@
 import psutil
 import tensorflow as tf
 import torch
+from jaxtyping import Int64, Shaped, UInt8
 
 from gigl.common import Uri
 from gigl.common.logger import Logger
@@ -19,10 +20,10 @@
 
 
 class LoadedEntityTensors(NamedTuple):
-    ids: torch.Tensor
-    features: Optional[torch.Tensor]
-    quantized_features: Optional[torch.Tensor]
-    labels: Optional[torch.Tensor]
+    ids: Union[Int64[torch.Tensor, "entities"], Int64[torch.Tensor, "2 entities"]]
+    features: Optional[Shaped[torch.Tensor, "entities feature_dim"]]
+    quantized_features: Optional[UInt8[torch.Tensor, "entities packed_feature_dim"]]
+    labels: Optional[Shaped[torch.Tensor, "entities labels"]]
 
 
 @dataclass(frozen=True)
@@ -443,9 +444,9 @@ def load_as_torch_tensors(
                 f"No files to load for rank: {self._rank} and entity type: {entity_type.name}, returning empty tensors."
             )
             empty_entity = (
-                torch.empty(0)
+                torch.empty(0, dtype=torch.int64)
                 if entity_type == FeatureTypes.NODE
-                else torch.empty(2, 0)
+                else torch.empty(2, 0, dtype=torch.int64)
             )
             if feature_keys:
                 empty_feature = torch.empty(0, serialized_tf_record_info.feature_dim)

gigl/common/data/export.py

@@ -20,6 +20,7 @@
 from google.cloud import bigquery
 from google.cloud.bigquery.job import LoadJob
 from google.cloud.exceptions import GoogleCloudError
+from jaxtyping import Int64, Shaped
 from typing_extensions import Self
 
 from gigl.common import GcsUri, LocalUri, Uri
@@ -249,10 +250,10 @@ def __init__(
 
     def add_embedding(
         self,
-        id_batch: torch.Tensor,
-        embedding_batch: torch.Tensor,
+        id_batch: Int64[torch.Tensor, "batch"],
+        embedding_batch: Shaped[torch.Tensor, "batch embedding_dim"],
         embedding_type: str,
-    ):
+    ) -> None:
         """
         Adds to the in-memory buffer the integer IDs and their corresponding embeddings.
 
@@ -308,10 +309,10 @@ def __init__(
 
     def add_prediction(
         self,
-        id_batch: torch.Tensor,
-        prediction_batch: torch.Tensor,
+        id_batch: Int64[torch.Tensor, "batch"],
+        prediction_batch: Shaped[torch.Tensor, "batch"],
         prediction_type: str,
-    ):
+    ) -> None:
         """
         Adds to the in-memory buffer the integer IDs and their corresponding predictions.
 

gigl/common/utils/feature_quantization/numpy_ops.py

@@ -6,17 +6,18 @@
 """
 
 import numpy as np
+from jaxtyping import Float, UInt8
 
 from gigl.common.utils.feature_quantization import SUPPORTED_QUANTIZATION_BITS
 
 
 def quantize_ndarray(
-    features: np.ndarray,
+    features: Float[np.ndarray, "entities feature_dim"],
     *,
     bits: int,
     clip_min: float | None = None,
     clip_max: float | None = None,
-) -> np.ndarray:
+) -> UInt8[np.ndarray, "entities packed_feature_dim"]:
     """Quantize a 2D float array into packed uint8 codes.
 
     For multi-bit quantization, `clip_min` and `clip_max` are required and

gigl/common/utils/feature_quantization/torch_ops.py

@@ -6,14 +6,15 @@
 """
 
 import torch
+from jaxtyping import Float32, UInt8
 
 from gigl.types.graph import FeatureQuantizationMetadata
 
 
 def dequantize_torch_tensor(
-    packed_features: torch.Tensor,
+    packed_features: UInt8[torch.Tensor, "... packed_feature_dim"],
     metadata: FeatureQuantizationMetadata,
-) -> torch.Tensor:
+) -> Float32[torch.Tensor, "... {metadata.quantized_feature_dim}"]:
     """Reconstruct approximate float features from packed uint8 codes."""
     q = metadata
 

gigl/distributed/base_sampler.py

@@ -16,6 +16,7 @@
 )
 from graphlearn_torch.typing import NodeType, as_str
 from graphlearn_torch.utils import reverse_edge_type
+from jaxtyping import Int64
 
 from gigl.common.logger import Logger
 from gigl.distributed.sampler import (
@@ -187,7 +188,7 @@ def _prepare_sample_loop_inputs(
     def _prepare_ablp_inputs(
         self,
         inputs: ABLPNodeSamplerInput,
-        input_seeds: torch.Tensor,
+        input_seeds: Int64[torch.Tensor, "{inputs.node.shape[0]}"],
         input_type: NodeType,
     ) -> SampleLoopInputs:
         """Prepare ABLP inputs with supervision nodes and label metadata.

gigl/distributed/dist_ablp_neighborloader.py

@@ -9,6 +9,7 @@
     MpDistSamplingWorkerOptions,
     RemoteDistSamplingWorkerOptions,
 )
+from jaxtyping import Int64
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.typing import EdgeType
 
@@ -74,8 +75,8 @@ def __init__(
         num_neighbors: Union[list[int], dict[EdgeType, list[int]]],
         input_nodes: Optional[
             Union[
-                torch.Tensor,
-                tuple[NodeType, torch.Tensor],
+                Int64[torch.Tensor, "nodes"],
+                tuple[NodeType, Int64[torch.Tensor, "nodes"]],
                 # Graph Store mode inputs
                 dict[int, ABLPInputNodes],
             ]
@@ -98,7 +99,7 @@ def __init__(
         local_process_world_size: Optional[int] = None,  # TODO: (svij) Deprecate this
         non_blocking_transfers: bool = True,
         use_label_edge_index_output: bool = False,
-    ):
+    ) -> None:
         """
         Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
 

gigl/distributed/dist_ppr_sampler.py

@@ -23,6 +23,7 @@
 )
 from graphlearn_torch.typing import EdgeType, NodeType
 from graphlearn_torch.utils import merge_dict, reverse_edge_type
+from jaxtyping import Int32
 
 from gigl.distributed.base_sampler import BaseDistNeighborSampler
 from gigl.distributed.utils.dist_typed_sampler import (
@@ -228,12 +229,15 @@ def __init__(
         max_ppr_nodes: int = 50,
         enable_residual_topup: bool = True,
         num_neighbors_per_hop: int = 100_000,
-        degree_tensors: Union[torch.Tensor, dict[NodeType, torch.Tensor]],
+        degree_tensors: Union[
+            Int32[torch.Tensor, "nodes"],
+            dict[NodeType, Int32[torch.Tensor, "_nodes"]],
+        ],
         max_fetch_iterations: Optional[int] = None,
         typed_channel_ratios: Optional[dict[TypedPPRChannelKey, float]] = None,
         include_sampled_edges: bool = False,
         **kwargs,
-    ):
+    ) -> None:
         super().__init__(*args, **kwargs)
         self._alpha = alpha
         if isinstance(max_ppr_nodes, bool) or max_ppr_nodes < 0: