diff --git a/CLAUDE.md b/CLAUDE.md index bd7d20026..8da4735b3 100644 --- a/CLAUDE.md +++ b/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 diff --git a/Makefile b/Makefile index ceddb35cd..318287829 100644 --- a/Makefile +++ b/Makefile @@ -248,7 +248,7 @@ 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" @@ -256,7 +256,7 @@ run_cora_nalp_e2e_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" @@ -264,7 +264,7 @@ run_cora_snc_e2e_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" @@ -272,7 +272,7 @@ run_cora_udl_e2e_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" @@ -280,7 +280,7 @@ run_dblp_nalp_e2e_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" @@ -288,7 +288,7 @@ run_hom_cora_sup_e2e_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" @@ -296,7 +296,7 @@ run_het_dblp_sup_e2e_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" @@ -304,7 +304,7 @@ run_hom_cora_sup_gs_e2e_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" @@ -312,7 +312,7 @@ run_het_dblp_sup_gs_e2e_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" @@ -320,7 +320,7 @@ run_hom_cora_snc_e2e_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" diff --git a/examples/link_prediction/graph_store/heterogeneous_training.py b/examples/link_prediction/graph_store/heterogeneous_training.py index 599178061..27d02672c 100644 --- a/examples/link_prediction/graph_store/heterogeneous_training.py +++ b/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: diff --git a/examples/link_prediction/graph_store/homogeneous_training.py b/examples/link_prediction/graph_store/homogeneous_training.py index a63e2bdc2..b32c74ac7 100644 --- a/examples/link_prediction/graph_store/homogeneous_training.py +++ b/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: diff --git a/examples/link_prediction/heterogeneous_training.py b/examples/link_prediction/heterogeneous_training.py index 1a6f66f44..3b91e93b3 100644 --- a/examples/link_prediction/heterogeneous_training.py +++ b/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: diff --git a/examples/link_prediction/homogeneous_training.py b/examples/link_prediction/homogeneous_training.py index 24814b456..85c416df2 100644 --- a/examples/link_prediction/homogeneous_training.py +++ b/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: diff --git a/examples/node_classification/models.py b/examples/node_classification/models.py index b55921963..2031b7bf8 100644 --- a/examples/node_classification/models.py +++ b/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. diff --git a/examples/tutorial/KDD_2025/heterogeneous_training.py b/examples/tutorial/KDD_2025/heterogeneous_training.py index d8a09646c..7035fff32 100644 --- a/examples/tutorial/KDD_2025/heterogeneous_training.py +++ b/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()]) diff --git a/gigl/common/data/dataloaders.py b/gigl/common/data/dataloaders.py index 67152c898..111cd309f 100644 --- a/gigl/common/data/dataloaders.py +++ b/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) diff --git a/gigl/common/data/export.py b/gigl/common/data/export.py index ba9af221b..0f15e18e3 100644 --- a/gigl/common/data/export.py +++ b/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. diff --git a/gigl/common/utils/feature_quantization/numpy_ops.py b/gigl/common/utils/feature_quantization/numpy_ops.py index e2fbd393a..be87c358f 100644 --- a/gigl/common/utils/feature_quantization/numpy_ops.py +++ b/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 diff --git a/gigl/common/utils/feature_quantization/torch_ops.py b/gigl/common/utils/feature_quantization/torch_ops.py index b91067584..49cddd30c 100644 --- a/gigl/common/utils/feature_quantization/torch_ops.py +++ b/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 diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 5601099e8..6c7b5a72b 100644 --- a/gigl/distributed/base_sampler.py +++ b/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. diff --git a/gigl/distributed/dist_ablp_neighborloader.py b/gigl/distributed/dist_ablp_neighborloader.py index 269334050..ab635b979 100644 --- a/gigl/distributed/dist_ablp_neighborloader.py +++ b/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. diff --git a/gigl/distributed/dist_ppr_sampler.py b/gigl/distributed/dist_ppr_sampler.py index e4d4d2da8..7ad1518c4 100644 --- a/gigl/distributed/dist_ppr_sampler.py +++ b/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: diff --git a/gigl/distributed/distributed_neighborloader.py b/gigl/distributed/distributed_neighborloader.py index 270bbfc45..7c00567b1 100644 --- a/gigl/distributed/distributed_neighborloader.py +++ b/gigl/distributed/distributed_neighborloader.py @@ -10,6 +10,7 @@ RemoteDistSamplingWorkerOptions, ) from graphlearn_torch.sampler import NodeSamplerInput +from jaxtyping import Int64 from torch_geometric.data import Data, HeteroData from torch_geometric.typing import EdgeType @@ -66,10 +67,10 @@ def __init__( num_neighbors: Union[list[int], dict[EdgeType, list[int]]], input_nodes: Optional[ Union[ - torch.Tensor, - Tuple[NodeType, torch.Tensor], - abc.Mapping[int, torch.Tensor], - Tuple[NodeType, abc.Mapping[int, torch.Tensor]], + Int64[torch.Tensor, "nodes"], + Tuple[NodeType, Int64[torch.Tensor, "nodes"]], + abc.Mapping[int, Int64[torch.Tensor, "nodes"]], + Tuple[NodeType, abc.Mapping[int, Int64[torch.Tensor, "nodes"]]], ] ] = None, num_workers: int = 1, @@ -88,7 +89,7 @@ def __init__( with_weight: bool = False, sampler_options: Optional[SamplerOptions] = None, non_blocking_transfers: bool = True, - ): + ) -> None: """ Distributed Neighbor Loader. Takes in some input nodes and samples neighbors from the dataset. diff --git a/gigl/distributed/sampler.py b/gigl/distributed/sampler.py index 1e01ee85f..9b101e67d 100644 --- a/gigl/distributed/sampler.py +++ b/gigl/distributed/sampler.py @@ -2,6 +2,7 @@ import torch from graphlearn_torch.sampler import NodeSamplerInput +from jaxtyping import Int64 from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.utils.share_memory import share_memory @@ -20,11 +21,15 @@ class ABLPNodeSamplerInput(NodeSamplerInput): def __init__( self, - node: torch.Tensor, + node: Int64[torch.Tensor, "anchors"], input_type: Optional[Union[str, NodeType]], - positive_label_by_edge_types: dict[EdgeType, torch.Tensor], - negative_label_by_edge_types: dict[EdgeType, torch.Tensor], - ): + positive_label_by_edge_types: dict[ + EdgeType, Int64[torch.Tensor, "anchors _positive_labels_per_anchor"] + ], + negative_label_by_edge_types: dict[ + EdgeType, Int64[torch.Tensor, "anchors _negative_labels_per_anchor"] + ], + ) -> None: """ Args: node (torch.Tensor): Anchor nodes to fanout from @@ -38,11 +43,15 @@ def __init__( self._negative_label_by_edge_types = negative_label_by_edge_types @property - def positive_label_by_edge_types(self) -> dict[EdgeType, torch.Tensor]: + def positive_label_by_edge_types( + self, + ) -> dict[EdgeType, Int64[torch.Tensor, "anchors _positive_labels_per_anchor"]]: return self._positive_label_by_edge_types @property - def negative_label_by_edge_types(self) -> dict[EdgeType, torch.Tensor]: + def negative_label_by_edge_types( + self, + ) -> dict[EdgeType, Int64[torch.Tensor, "anchors _negative_labels_per_anchor"]]: return self._negative_label_by_edge_types def __len__(self) -> int: diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 5868e6394..d463cb4d5 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -12,12 +12,13 @@ """ import math -from typing import Callable, Literal, Optional, cast +from typing import Callable, Literal, Optional, Union, cast import torch import torch.nn as nn import torch.nn.functional as F import torch_geometric.data.hetero_data +from jaxtyping import Bool, Float, Int64 from torch import Tensor from gigl.src.common.types.graph_data import EdgeType, NodeType @@ -189,7 +190,9 @@ def reset_parameters(self) -> None: if layer.bias is not None: nn.init.zeros_(layer.bias) - def forward(self, x: Tensor) -> Tensor: + def forward( + self, x: Float[Tensor, "batch sequence model_dim"] + ) -> Float[Tensor, "batch sequence model_dim"]: """Forward pass. Args: @@ -442,12 +445,14 @@ def _reset_relation_message_attention_parameters(self) -> None: def forward( self, - x: Tensor, - attn_bias: Optional[Tensor] = None, - valid_mask: Optional[Tensor] = None, - pairwise_relation_indices: Optional[Tensor] = None, + x: Float[Tensor, "batch sequence model_dim"], + attn_bias: Optional[ + Union[Float[Tensor, ""], Float[Tensor, "... #sequence"]] + ] = None, + valid_mask: Optional[Bool[Tensor, "batch sequence"]] = None, + pairwise_relation_indices: Optional[Int64[Tensor, "relation_edges 4"]] = None, query_seq_len: Optional[int] = None, - ) -> Tensor: + ) -> Float[Tensor, "batch query_sequence model_dim"]: """Forward pass. Args: @@ -1480,9 +1485,9 @@ def forward( self, data: torch_geometric.data.hetero_data.HeteroData, anchor_node_type: Optional[NodeType] = None, - anchor_node_ids: Optional[Tensor] = None, + anchor_node_ids: Optional[Int64[Tensor, "anchors"]] = None, device: Optional[torch.device] = None, - ) -> torch.Tensor: + ) -> Float[torch.Tensor, "anchors output_dim"]: """Run the forward pass of the Graph Transformer encoder. Args: diff --git a/gigl/nn/loss.py b/gigl/nn/loss.py index f0e11247c..91fc7c179 100644 --- a/gigl/nn/loss.py +++ b/gigl/nn/loss.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn +from jaxtyping import Float, Int64 class RetrievalLoss(nn.Module): @@ -38,12 +39,14 @@ def __init__( def _calculate_batch_retrieval_loss( self, - scores: torch.Tensor, - candidate_sampling_probability: Optional[torch.Tensor] = None, - query_ids: Optional[torch.Tensor] = None, - candidate_ids: Optional[torch.Tensor] = None, + scores: Float[torch.Tensor, "queries candidates"], + candidate_sampling_probability: Optional[ + Float[torch.Tensor, "candidates"] + ] = None, + query_ids: Optional[Int64[torch.Tensor, "queries"]] = None, + candidate_ids: Optional[Int64[torch.Tensor, "candidates"]] = None, device: torch.device = torch.device("cpu"), - ) -> torch.Tensor: + ) -> Float[torch.Tensor, ""]: """ Args: scores: [num_queries, num_candidates] tensor of candidate and query embeddings similarity @@ -160,12 +163,14 @@ def _mask_by_candidate_ids( def forward( self, - repeated_candidate_scores: torch.Tensor, - candidate_ids: torch.Tensor, - repeated_query_ids: torch.Tensor, + repeated_candidate_scores: Float[torch.Tensor, "queries candidates"], + candidate_ids: Int64[torch.Tensor, "candidates"], + repeated_query_ids: Int64[torch.Tensor, "queries"], device: torch.device, - candidate_sampling_probability: Optional[torch.Tensor] = None, - ): + candidate_sampling_probability: Optional[ + Float[torch.Tensor, "candidates"] + ] = None, + ) -> Float[torch.Tensor, ""]: """ Args: repeated_candidate_scores (torch.Tensor): The prediction scores between each repeated query users and each candidates. In this case, `repeated` means diff --git a/gigl/nn/models.py b/gigl/nn/models.py index 6016561c3..e593013a7 100644 --- a/gigl/nn/models.py +++ b/gigl/nn/models.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn +from jaxtyping import Float, Int64 from torch.nn.parallel import DistributedDataParallel from torch_geometric.data import Data, HeteroData from torch_geometric.nn.conv import LGConv @@ -38,7 +39,10 @@ def forward( data: Union[Data, HeteroData], device: torch.device, output_node_types: Optional[list[NodeType]] = None, - ) -> Union[torch.Tensor, dict[NodeType, torch.Tensor]]: + ) -> Union[ + Float[torch.Tensor, "_nodes embedding_dim"], + dict[NodeType, Float[torch.Tensor, "_nodes embedding_dim"]], + ]: if isinstance(data, HeteroData): if output_node_types is None: raise ValueError( @@ -52,9 +56,9 @@ def forward( def decode( self, - query_embeddings: torch.Tensor, - candidate_embeddings: torch.Tensor, - ) -> torch.Tensor: + query_embeddings: Float[torch.Tensor, "queries embedding_dim"], + candidate_embeddings: Float[torch.Tensor, "candidates embedding_dim"], + ) -> Float[torch.Tensor, "queries candidates"]: return self._decoder( query_embeddings=query_embeddings, candidate_embeddings=candidate_embeddings, @@ -216,8 +220,11 @@ def forward( data: Union[Data, HeteroData], device: torch.device, output_node_types: Optional[list[NodeType]] = None, - anchor_node_ids: Optional[torch.Tensor] = None, - ) -> Union[torch.Tensor, dict[NodeType, torch.Tensor]]: + anchor_node_ids: Optional[Int64[torch.Tensor, "anchors"]] = None, + ) -> Union[ + Float[torch.Tensor, "_nodes embedding_dim"], + dict[NodeType, Float[torch.Tensor, "_nodes embedding_dim"]], + ]: """ Forward pass of the LightGCN model. @@ -325,7 +332,7 @@ def _forward_homogeneous( def _lookup_embeddings_for_single_node_type( self, node_type: str, ids: torch.Tensor - ) -> torch.Tensor: + ) -> Union[torch.Tensor, Awaitable[torch.Tensor]]: """ Fetch per-ID embeddings for a single node type using EmbeddingBagCollection. @@ -340,7 +347,9 @@ def _lookup_embeddings_for_single_node_type( ids (torch.Tensor): Node IDs to look up, shape [batch_size]. Returns: - torch.Tensor: Embeddings for the requested node type, shape [batch_size, embedding_dim]. + Embeddings for the requested node type, shape + ``[batch_size, embedding_dim]``. Distributed embedding collections + return an awaitable that resolves to the tensor. """ if node_type not in self._feature_keys: raise KeyError( diff --git a/gigl/src/common/modeling_task_specs/utils/infer.py b/gigl/src/common/modeling_task_specs/utils/infer.py index a5adb18d3..6cc418595 100644 --- a/gigl/src/common/modeling_task_specs/utils/infer.py +++ b/gigl/src/common/modeling_task_specs/utils/infer.py @@ -1,8 +1,9 @@ from collections import defaultdict -from typing import Set, Union +from typing import Set, Union, cast import torch import torch.nn as nn +from jaxtyping import Float, Int64 from torch_geometric.data import Data from torch_geometric.data.hetero_data import HeteroData @@ -73,10 +74,10 @@ def infer_training_batch( def infer_root_embeddings( model: Union[torch.nn.parallel.DistributedDataParallel, nn.Module], graph: Union[Data, HeteroData], - root_node_indices: torch.LongTensor, + root_node_indices: Int64[torch.LongTensor, "roots"], gbml_config_pb_wrapper: GbmlConfigPbWrapper, device: torch.device, -) -> torch.Tensor: +) -> Float[torch.FloatTensor, "roots embedding_dim"]: batch_graph = graph.to(device=device) batch_root_node_indices = root_node_indices.to(device=device) output_node_types = list( @@ -217,18 +218,27 @@ def infer_task_inputs( condensed_node_type ].to(device=device) ) - random_neg_root_embeddings[condensed_node_type] = ( - random_neg_embeddings[condensed_node_type][random_neg_root_node_indices] - if random_neg_root_node_indices.numel() - else torch.FloatTensor([]).to(device=device) # ty: ignore[invalid-assignment] TODO(ty-torch-tensor-specialization): fix ty Tensor vs FloatTensor/LongTensor specialization. + random_neg_root_embeddings[condensed_node_type] = cast( + torch.FloatTensor, + ( + random_neg_embeddings[condensed_node_type][random_neg_root_node_indices] + if random_neg_root_node_indices.numel() + else random_neg_embeddings[condensed_node_type].new_empty( + (0, random_neg_embeddings[condensed_node_type].shape[-1]) + ) + ), ) if ModelResultType.batch_scores in batch_result_types or should_eval: - random_neg_scores[condensed_node_type] = ( - decoder( - query_embeddings, random_neg_root_embeddings[condensed_node_type] - ) - if random_neg_root_embeddings[condensed_node_type].numel() - else torch.FloatTensor([]).to(device=device) # ty: ignore[invalid-assignment] TODO(ty-torch-tensor-specialization): fix ty Tensor vs FloatTensor/LongTensor specialization. + random_neg_scores[condensed_node_type] = cast( + torch.FloatTensor, + ( + decoder( + query_embeddings, + random_neg_root_embeddings[condensed_node_type], + ) + if random_neg_root_embeddings[condensed_node_type].numel() + else query_embeddings.new_empty((query_embeddings.shape[0], 0)) + ), ) # Loop through all root nodes and populate ids, embeddings, and scores per condensed edge type @@ -285,7 +295,7 @@ def infer_task_inputs( ], ) if pos_nodes.numel() - else torch.FloatTensor([]).to(device=device) + else query_embeddings.new_empty((1, 0)) ) hard_neg_scores = ( decoder( @@ -295,7 +305,7 @@ def infer_task_inputs( ], ) if hard_neg_nodes.numel() - else torch.FloatTensor([]).to(device=device) + else query_embeddings.new_empty((1, 0)) ) random_neg_scores_root = random_neg_scores[ condensed_supervision_target_node_type @@ -322,15 +332,35 @@ def infer_task_inputs( ) = gbml_config_pb_wrapper.graph_metadata_pb_wrapper.condensed_edge_type_to_condensed_node_types[ condensed_supervision_edge_type ] - pos_embeddings[condensed_supervision_edge_type] = ( - torch.cat(tuple(_pos_embeddings[condensed_supervision_edge_type])) - if len(_pos_embeddings[condensed_supervision_edge_type]) - else torch.tensor([]) # ty: ignore[invalid-assignment] TODO(ty-torch-tensor-specialization): fix ty Tensor vs FloatTensor/LongTensor specialization. + pos_embeddings[condensed_supervision_edge_type] = cast( + torch.FloatTensor, + ( + torch.cat(tuple(_pos_embeddings[condensed_supervision_edge_type])) + if len(_pos_embeddings[condensed_supervision_edge_type]) + else main_embeddings[condensed_supervision_target_node_type].new_empty( + ( + 0, + main_embeddings[condensed_supervision_target_node_type].shape[ + -1 + ], + ) + ) + ), ) - hard_neg_embeddings[condensed_supervision_edge_type] = ( - torch.cat(tuple(_hard_neg_embeddings[condensed_supervision_edge_type])) - if len(_hard_neg_embeddings[condensed_supervision_edge_type]) - else torch.tensor([]) # ty: ignore[invalid-assignment] TODO(ty-torch-tensor-specialization): fix ty Tensor vs FloatTensor/LongTensor specialization. + hard_neg_embeddings[condensed_supervision_edge_type] = cast( + torch.FloatTensor, + ( + torch.cat(tuple(_hard_neg_embeddings[condensed_supervision_edge_type])) + if len(_hard_neg_embeddings[condensed_supervision_edge_type]) + else main_embeddings[condensed_supervision_target_node_type].new_empty( + ( + 0, + main_embeddings[condensed_supervision_target_node_type].shape[ + -1 + ], + ) + ) + ), ) repeated_anchor_embeddings[condensed_supervision_edge_type] = ( @@ -430,7 +460,7 @@ def infer_task_inputs( candidate_embeddings, ) if repeated_anchor_embeddings[condensed_supervision_edge_type].numel() - else torch.tensor([]) + else candidate_embeddings.new_empty((0, candidate_embeddings.shape[0])) ) batch_combined_scores[condensed_supervision_edge_type] = ( diff --git a/gigl/src/common/models/layers/count_min_sketch.py b/gigl/src/common/models/layers/count_min_sketch.py index 2fffe129f..48f4aacbd 100644 --- a/gigl/src/common/models/layers/count_min_sketch.py +++ b/gigl/src/common/models/layers/count_min_sketch.py @@ -2,6 +2,7 @@ import numpy as np import torch +from jaxtyping import Float32, Int64 from gigl.common.logger import Logger @@ -55,7 +56,7 @@ def add(self, item: Any, delta: int = 1) -> None: self.__table[i][hashed_value % self.__width] += delta self.__total += delta - def add_torch_long_tensor(self, tensor: torch.LongTensor) -> None: + def add_torch_long_tensor(self, tensor: Int64[torch.LongTensor, "items"]) -> None: """ Add all items in a torch long tensor to the sketch """ @@ -79,7 +80,9 @@ def estimate(self, item: Any) -> int: for i, hashed_value in enumerate(hashed_values) ) - def estimate_torch_long_tensor(self, tensor: torch.LongTensor) -> torch.LongTensor: + def estimate_torch_long_tensor( + self, tensor: Int64[torch.LongTensor, "items"] + ) -> Int64[torch.LongTensor, "items"]: """ Return the estimated count of all items in a torch long tensor """ @@ -97,8 +100,10 @@ def get_table(self) -> np.ndarray: def calculate_in_batch_candidate_sampling_probability( - frequency_tensor: torch.LongTensor, total_cnt: int, batch_size: int -) -> torch.Tensor: + frequency_tensor: Int64[torch.LongTensor, "candidates"], + total_cnt: int, + batch_size: int, +) -> Float32[torch.FloatTensor, "candidates"]: """ Calculate in batch negative sampling rate given the frequency tensor, total count and batch size. Please see https://www.tensorflow.org/extras/candidate_sampling.pdf for more details @@ -117,4 +122,4 @@ def calculate_in_batch_candidate_sampling_probability( estimated_prob: torch.FloatTensor = ( batch_size * frequency_tensor.float() / total_cnt ) # ty: ignore[invalid-assignment] TODO(ty-torch-tensor-specialization): fix ty Tensor vs FloatTensor/LongTensor specialization. - return estimated_prob.clamp(max=1.0) + return estimated_prob.clamp(max=1.0) # ty: ignore[invalid-return-type] TODO(ty-torch-tensor-specialization): fix ty Tensor vs FloatTensor/LongTensor specialization. diff --git a/gigl/src/common/models/layers/decoder.py b/gigl/src/common/models/layers/decoder.py index 381c1c600..dd0ef2a6e 100644 --- a/gigl/src/common/models/layers/decoder.py +++ b/gigl/src/common/models/layers/decoder.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from jaxtyping import Float from torch_geometric.nn.models import MLP @@ -61,7 +62,11 @@ def __init__( norm=norm, ) - def forward(self, query_embeddings, candidate_embeddings) -> torch.Tensor: + def forward( + self, + query_embeddings: Float[torch.Tensor, "queries embedding_dim"], + candidate_embeddings: Float[torch.Tensor, "candidates embedding_dim"], + ) -> Float[torch.Tensor, "queries candidates"]: if self.decoder_type.value == "inner_product": scores = torch.mm(query_embeddings, candidate_embeddings.T) elif self.decoder_type.value == "hadamard_MLP": diff --git a/gigl/src/common/models/layers/feature_interaction.py b/gigl/src/common/models/layers/feature_interaction.py index f0ccf2126..fa7e7234f 100644 --- a/gigl/src/common/models/layers/feature_interaction.py +++ b/gigl/src/common/models/layers/feature_interaction.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn +from jaxtyping import Float class DCNCross(nn.Module): @@ -60,8 +61,10 @@ def __init__( ) def forward( - self, x0: torch.Tensor, x: Optional[torch.Tensor] = None - ) -> torch.Tensor: + self, + x0: Float[torch.Tensor, "... feature_dim"], + x: Optional[Float[torch.Tensor, "... feature_dim"]] = None, + ) -> Float[torch.Tensor, "... feature_dim"]: """ Computes the feature cross. Args: @@ -141,7 +144,9 @@ def __init__( ) ) - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward( + self, x: Float[torch.Tensor, "... feature_dim"] + ) -> Float[torch.Tensor, "... feature_dim"]: x0, xl = x, x for i in range(self._num_layers): xl = self._layers[i](x0, xl) diff --git a/gigl/src/common/models/layers/loss.py b/gigl/src/common/models/layers/loss.py index 9ca934126..5ab5299b7 100644 --- a/gigl/src/common/models/layers/loss.py +++ b/gigl/src/common/models/layers/loss.py @@ -6,6 +6,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from jaxtyping import Float, Int64 from gigl.common.logger import Logger from gigl.src.common.types.graph_data import CondensedEdgeType @@ -42,11 +43,11 @@ def __init__( def _calculate_margin_loss( self, - pos_scores: torch.Tensor, - hard_neg_scores: torch.Tensor, - random_neg_scores: torch.Tensor, + pos_scores: Float[torch.Tensor, "1 positives"], + hard_neg_scores: Float[torch.Tensor, "1 hard_negatives"], + random_neg_scores: Float[torch.Tensor, "1 random_negatives"], device: torch.device = torch.device("cpu"), - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: all_neg_scores = torch.cat( (hard_neg_scores, random_neg_scores), dim=1, @@ -75,7 +76,7 @@ def forward( self, loss_input: list[dict[CondensedEdgeType, BatchScores]], device: torch.device = torch.device("cpu"), - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: batch_loss = torch.tensor(0.0).to(device=device) batch_size = 0 # In case we have an empty list as input, avoids division by zero error @@ -127,11 +128,11 @@ def __init__( def _calculate_softmax_loss( self, - pos_scores: torch.Tensor, - hard_neg_scores: torch.Tensor, - random_neg_scores: torch.Tensor, + pos_scores: Float[torch.Tensor, "1 positives"], + hard_neg_scores: Float[torch.Tensor, "1 hard_negatives"], + random_neg_scores: Float[torch.Tensor, "1 random_negatives"], device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: all_neg_scores = torch.cat( (hard_neg_scores, random_neg_scores), dim=1, @@ -162,7 +163,7 @@ def forward( self, loss_input: list[dict[CondensedEdgeType, BatchScores]], device: torch.device = torch.device("cpu"), - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: batch_loss = torch.tensor(0.0).to(device=device) batch_size = 0 # In case we have an empty list as input, avoids division by zero error @@ -224,12 +225,14 @@ def __init__( def calculate_batch_retrieval_loss( self, - scores: torch.Tensor, - candidate_sampling_probability: Optional[torch.Tensor] = None, - query_ids: Optional[torch.Tensor] = None, - candidate_ids: Optional[torch.Tensor] = None, + scores: Float[torch.Tensor, "queries candidates"], + candidate_sampling_probability: Optional[ + Float[torch.Tensor, "candidates"] + ] = None, + query_ids: Optional[Int64[torch.Tensor, "queries"]] = None, + candidate_ids: Optional[Int64[torch.Tensor, "candidates"]] = None, device: torch.device = torch.device("cpu"), - ) -> torch.Tensor: + ) -> Float[torch.Tensor, ""]: """ Args: scores: [num_queries, num_candidates] tensor of candidate and query embeddings similarity @@ -347,10 +350,12 @@ def _mask_by_candidate_ids( def forward( self, batch_combined_scores: BatchCombinedScores, - repeated_query_embeddings: torch.FloatTensor, - candidate_sampling_probability: Optional[torch.FloatTensor] = None, + repeated_query_embeddings: Float[torch.Tensor, "queries embedding_dim"], + candidate_sampling_probability: Optional[ + Float[torch.Tensor, "candidates"] + ] = None, device: torch.device = torch.device("cpu"), - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: candidate_ids = torch.cat( ( batch_combined_scores.positive_ids.to(device=device), @@ -389,10 +394,10 @@ def __init__( def forward( self, - h1: torch.Tensor, - h2: torch.Tensor, + h1: Float[torch.Tensor, "nodes embedding_dim"], + h2: Float[torch.Tensor, "nodes embedding_dim"], device: torch.device = torch.device("cpu"), - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: """ Args: h1 (torch.Tensor): First input tensor @@ -445,9 +450,9 @@ def __init__( def forward( self, - x_target: torch.Tensor, - x_pred: torch.Tensor, - ) -> Tuple[torch.Tensor, int]: + x_target: Float[torch.Tensor, "nodes feature_dim"], + x_pred: Float[torch.Tensor, "nodes feature_dim"], + ) -> Tuple[Float[torch.Tensor, ""], int]: x = F.normalize(x_target, p=2, dim=-1) # SCE Loss Computation y = F.normalize(x_pred, p=2, dim=-1) loss = (1 - (x * y).sum(dim=-1)).pow_(self.alpha) @@ -470,11 +475,11 @@ def __init__( def forward( self, - h1: torch.Tensor, - h2: torch.Tensor, + h1: Float[torch.Tensor, "nodes embedding_dim"], + h2: Float[torch.Tensor, "nodes embedding_dim"], N: int, device: torch.device = torch.device("cpu"), - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: """ Args: h1 (torch.Tensor): First input tensor @@ -517,10 +522,10 @@ def __init__( def forward( self, - z_a: torch.Tensor, - z_b: torch.Tensor, + z_a: Float[torch.Tensor, "nodes feature_dim"], + z_b: Float[torch.Tensor, "nodes feature_dim"], device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: """ Args: z_a (torch.Tensor): First input matrix @@ -557,11 +562,11 @@ class BGRLLoss(nn.Module): def forward( self, - q1: torch.Tensor, - q2: torch.Tensor, - y1: torch.Tensor, - y2: torch.Tensor, - ) -> Tuple[torch.Tensor, int]: + q1: Float[torch.Tensor, "nodes embedding_dim"], + q2: Float[torch.Tensor, "nodes embedding_dim"], + y1: Float[torch.Tensor, "nodes embedding_dim"], + y2: Float[torch.Tensor, "nodes embedding_dim"], + ) -> Tuple[Float[torch.Tensor, ""], int]: loss = ( 2 - F.cosine_similarity(q1, y2.detach(), dim=-1).mean() @@ -586,12 +591,12 @@ def __init__( def forward( self, - q1: torch.Tensor, - q2: torch.Tensor, - y1: torch.Tensor, - y2: torch.Tensor, - neg_y: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, int]: + q1: Float[torch.Tensor, "nodes embedding_dim"], + q2: Float[torch.Tensor, "nodes embedding_dim"], + y1: Float[torch.Tensor, "nodes embedding_dim"], + y2: Float[torch.Tensor, "nodes embedding_dim"], + neg_y: Optional[Float[torch.Tensor, "nodes embedding_dim"]], + ) -> Tuple[Float[torch.Tensor, ""], int]: sim1 = F.cosine_similarity(q1, y2.detach()).mean() sim2 = F.cosine_similarity(q2, y1.detach()).mean() neg_sim1 = F.cosine_similarity(q1, neg_y.detach()).mean() # type: ignore @@ -615,8 +620,10 @@ def __init__( self.alpha = alpha def forward( - self, user_embeddings: torch.Tensor, item_embeddings: torch.Tensor - ) -> torch.Tensor: + self, + user_embeddings: Float[torch.Tensor, "pairs embedding_dim"], + item_embeddings: Float[torch.Tensor, "pairs embedding_dim"], + ) -> Float[torch.Tensor, ""]: return ( (user_embeddings - item_embeddings).norm(p=2, dim=1).pow(self.alpha).mean() ) @@ -635,8 +642,10 @@ def __init__( self.temperature = temperature def forward( - self, user_embeddings: torch.Tensor, item_embeddings: torch.Tensor - ) -> torch.Tensor: + self, + user_embeddings: Float[torch.Tensor, "_users embedding_dim"], + item_embeddings: Float[torch.Tensor, "_items embedding_dim"], + ) -> Float[torch.Tensor, ""]: user_uniformity = ( torch.pdist(user_embeddings, p=2) .pow(2) @@ -672,9 +681,9 @@ def __init__( def forward( self, - student_scores: torch.Tensor, - teacher_scores: torch.Tensor, - ) -> torch.Tensor: + student_scores: Float[torch.Tensor, "batch classes"], + teacher_scores: Float[torch.Tensor, "batch classes"], + ) -> Float[torch.Tensor, ""]: y_s = F.log_softmax(student_scores / self.kl_temperature, dim=-1) y_t = F.softmax(teacher_scores / self.kl_temperature, dim=-1) loss = ( @@ -704,10 +713,10 @@ def __init__( def forward( self, - student_scores: torch.Tensor, - teacher_scores: torch.Tensor, + student_scores: Float[torch.Tensor, "batch candidates"], + teacher_scores: Float[torch.Tensor, "batch candidates"], device: torch.device, - ) -> torch.Tensor: + ) -> Float[torch.Tensor, ""]: dim_pairs = [ x for x in itertools.combinations(range(student_scores.shape[1]), r=2) ] diff --git a/gigl/src/common/models/layers/task.py b/gigl/src/common/models/layers/task.py index 88515675a..cfae0baeb 100644 --- a/gigl/src/common/models/layers/task.py +++ b/gigl/src/common/models/layers/task.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn +from jaxtyping import Float from torch_geometric.nn import GraphConv from gigl.common.logger import Logger @@ -46,7 +47,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: raise NotImplementedError @property @@ -73,7 +74,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: assert len(task_input.batch_scores) > 0 return self.loss(loss_input=task_input.batch_scores, device=device) @@ -96,7 +97,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: assert len(task_input.batch_scores) > 0 return self.loss(loss_input=task_input.batch_scores, device=device) @@ -143,7 +144,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: assert len(task_input.batch_combined_scores) > 0 assert task_input.batch_embeddings is not None running_loss = torch.tensor(0.0, device=device) @@ -244,7 +245,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: main_batch = task_input.input_batch.main_batch augmented_graph_1 = get_augmented_graph( graph=main_batch.graph.to(device=device), @@ -306,7 +307,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: # TODO (mkolodner) Update GraphMAE logic to work in both heterogeneous use case if gbml_config_pb_wrapper.graph_metadata_pb_wrapper.is_heterogeneous: raise NotImplementedError( @@ -383,7 +384,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: main_batch = task_input.input_batch.main_batch augmented_graph_1 = get_augmented_graph( graph=main_batch.graph.to(device=device), @@ -445,7 +446,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: main_batch = task_input.input_batch.main_batch augmented_graph_1 = get_augmented_graph( graph=main_batch.graph.to(device=device), @@ -517,7 +518,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: main_batch = task_input.input_batch.main_batch augmented_graph_1 = get_augmented_graph( graph=main_batch.graph.to(device=device), @@ -603,7 +604,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: main_batch = task_input.input_batch.main_batch augmented_graph_1 = get_augmented_graph( graph=main_batch.graph.to(device=device), @@ -674,7 +675,7 @@ def forward( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, int]: + ) -> Tuple[Float[torch.Tensor, ""], int]: assert task_input.batch_embeddings is not None batch_embeddings = task_input.batch_embeddings running_loss = torch.tensor(0.0, device=device) @@ -728,7 +729,7 @@ def calculate_losses( gbml_config_pb_wrapper: GbmlConfigPbWrapper, should_eval: bool, device: torch.device, - ) -> Tuple[torch.Tensor, dict[str, float]]: + ) -> Tuple[Float[torch.Tensor, ""], dict[str, float]]: loss_to_val_map: dict[str, float] = {} loss_to_batch_size_map: dict[str, int] = {} for task, weight in self._get_all_tasks(): @@ -744,7 +745,7 @@ def calculate_losses( for loss_type in loss_to_val_map: cur_loss = loss_to_val_map[loss_type] sample_wise_loss += cur_loss / loss_to_batch_size_map[loss_type] - final_loss: torch.Tensor + final_loss: Float[torch.Tensor, ""] final_loss_map: dict[str, float] final_loss = sample_wise_loss diff --git a/gigl/src/common/models/pyg/heterogeneous.py b/gigl/src/common/models/pyg/heterogeneous.py index 8fb5cc18a..bca18de82 100644 --- a/gigl/src/common/models/pyg/heterogeneous.py +++ b/gigl/src/common/models/pyg/heterogeneous.py @@ -2,6 +2,7 @@ import torch import torch_geometric.data +from jaxtyping import Float from torch import nn from torch.nn import functional as F from torch_geometric.nn import Linear @@ -73,7 +74,7 @@ def forward( data: torch_geometric.data.hetero_data.HeteroData, output_node_types: list[NodeType], device: torch.device, - ) -> dict[NodeType, torch.Tensor]: + ) -> dict[NodeType, Float[torch.Tensor, "_nodes output_dim"]]: """ Runs the forward pass of the module Args: @@ -221,7 +222,7 @@ def forward( data: torch_geometric.data.hetero_data.HeteroData, output_node_types: list[NodeType], device: torch.device, - ) -> dict[NodeType, torch.Tensor]: + ) -> dict[NodeType, Float[torch.Tensor, "_nodes output_dim"]]: # Align dimensions across all node-types and all edge-types, resp. x_dict = { node_type: self.node_type_lin_dict[node_type](x) diff --git a/gigl/src/common/models/pyg/homogeneous.py b/gigl/src/common/models/pyg/homogeneous.py index 5af61c82b..6f21f7270 100644 --- a/gigl/src/common/models/pyg/homogeneous.py +++ b/gigl/src/common/models/pyg/homogeneous.py @@ -4,6 +4,7 @@ import torch.nn as nn import torch.nn.functional as F import torch_geometric.data +from jaxtyping import Float from torch_geometric.nn import ( GATConv, GATv2Conv, @@ -108,7 +109,7 @@ def forward( self, data: torch_geometric.data.Data, device: Optional[torch.device] = None, - ) -> torch.Tensor: + ) -> Float[torch.Tensor, "nodes output_dim"]: x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr # pass selected features through an embedding layer if self.feature_embedding_layer: @@ -531,7 +532,9 @@ def __init__( in_channels=hid_dim, out_channels=out_dim, **remaining_kwargs ) - def forward(self, data: torch_geometric.data.Data) -> torch.Tensor: + def forward( + self, data: torch_geometric.data.Data + ) -> Float[torch.Tensor, "nodes output_dim"]: x, edge_index = data.x, data.edge_index x = self.conv1(x, edge_index) x = F.relu(x) diff --git a/gigl/src/common/models/pyg/link_prediction.py b/gigl/src/common/models/pyg/link_prediction.py index bc266b41b..32f8f7346 100644 --- a/gigl/src/common/models/pyg/link_prediction.py +++ b/gigl/src/common/models/pyg/link_prediction.py @@ -3,6 +3,7 @@ import torch import torch.nn as nn import torch_geometric +from jaxtyping import Float from gigl.common.logger import Logger from gigl.src.common.models.layers.decoder import LinkPredictionDecoder @@ -45,7 +46,7 @@ def forward( ], output_node_types: list[NodeType], device: torch.device, - ) -> dict[NodeType, torch.Tensor]: + ) -> dict[NodeType, Float[torch.Tensor, "_nodes embedding_dim"]]: if isinstance(data, torch_geometric.data.hetero_data.HeteroData): return self.__encoder( data=data, output_node_types=output_node_types, device=device @@ -59,9 +60,9 @@ def forward( def decode( self, - query_embeddings: torch.Tensor, - candidate_embeddings: torch.Tensor, - ) -> torch.Tensor: + query_embeddings: Float[torch.Tensor, "queries embedding_dim"], + candidate_embeddings: Float[torch.Tensor, "candidates embedding_dim"], + ) -> Float[torch.Tensor, "queries candidates"]: return self.__decoder( query_embeddings=query_embeddings, candidate_embeddings=candidate_embeddings, diff --git a/gigl/src/common/models/pyg/nn/models/feature_embedding.py b/gigl/src/common/models/pyg/nn/models/feature_embedding.py index 1ddde48d6..8bb813f30 100644 --- a/gigl/src/common/models/pyg/nn/models/feature_embedding.py +++ b/gigl/src/common/models/pyg/nn/models/feature_embedding.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn +from jaxtyping import Float from tensorflow_metadata.proto.v0.schema_pb2 import Feature from gigl.common.logger import Logger @@ -122,7 +123,9 @@ def __init__( ) self.__out_dim += emb_dim - feat_dim # adjust out_dim based on emb_dim - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward( + self, x: Float[torch.Tensor, "nodes input_dim"] + ) -> Float[torch.Tensor, "nodes output_dim"]: x_non_embed = filter_features( feature_schema=self.__feature_schema, feature_names=list(self.__non_embed_features), diff --git a/gigl/src/common/models/pyg/nn/models/feature_interaction.py b/gigl/src/common/models/pyg/nn/models/feature_interaction.py index 850b6523d..77028e208 100644 --- a/gigl/src/common/models/pyg/nn/models/feature_interaction.py +++ b/gigl/src/common/models/pyg/nn/models/feature_interaction.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from jaxtyping import Float from torch_geometric.nn.models import MLP from gigl.src.common.models.layers.feature_interaction import DCNv2 @@ -60,7 +61,9 @@ def __init__( **mlp_feats_kwargs, ) - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward( + self, x: Float[torch.Tensor, "nodes input_dim"] + ) -> Float[torch.Tensor, "nodes output_dim"]: if self.combination_mode and self.combination_mode == CombinationMode.parallel: assert isinstance(self.dcnv2, nn.Module) and isinstance(self.mlp, nn.Module) x_cross = self.dcnv2(x) diff --git a/gigl/src/common/models/pyg/nn/models/jumping_knowledge.py b/gigl/src/common/models/pyg/nn/models/jumping_knowledge.py index 1d5f0403e..b883a8fcc 100644 --- a/gigl/src/common/models/pyg/nn/models/jumping_knowledge.py +++ b/gigl/src/common/models/pyg/nn/models/jumping_knowledge.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn +from jaxtyping import Float from torch import Tensor from torch_geometric.nn.conv import MessagePassing from torch_geometric.typing import OptPairTensor # noqa @@ -99,7 +100,9 @@ def reset_parameters(self): self.att.reset_parameters() self.output_linear.reset_parameters() - def forward(self, xs: list[torch.Tensor]) -> Tensor: + def forward( + self, xs: list[Float[torch.Tensor, "nodes hidden_dim"]] + ) -> Float[Tensor, "nodes output_dim"]: r""" Args: xs (list[torch.Tensor]): List containing the layer-wise diff --git a/gigl/src/common/types/task_inputs.py b/gigl/src/common/types/task_inputs.py index 15a0f420d..e6925dc4e 100644 --- a/gigl/src/common/types/task_inputs.py +++ b/gigl/src/common/types/task_inputs.py @@ -2,6 +2,7 @@ from typing import Optional import torch +from jaxtyping import Float32, Int64 from gigl.src.common.types.graph_data import CondensedEdgeType, CondensedNodeType from gigl.src.training.v1.lib.data_loaders.node_anchor_based_link_prediction_data_loader import ( @@ -22,29 +23,38 @@ class InputBatch: # Returns the embeddings after being forward through encoder model @dataclass class BatchEmbeddings: - query_embeddings: torch.FloatTensor - repeated_query_embeddings: dict[CondensedEdgeType, torch.FloatTensor] - pos_embeddings: dict[CondensedEdgeType, torch.FloatTensor] - hard_neg_embeddings: dict[CondensedEdgeType, torch.FloatTensor] - random_neg_embeddings: dict[CondensedNodeType, torch.FloatTensor] + query_embeddings: Float32[torch.FloatTensor, "queries embedding_dim"] + repeated_query_embeddings: dict[ + CondensedEdgeType, Float32[torch.FloatTensor, "_queries embedding_dim"] + ] + pos_embeddings: dict[ + CondensedEdgeType, Float32[torch.FloatTensor, "_positives embedding_dim"] + ] + hard_neg_embeddings: dict[ + CondensedEdgeType, Float32[torch.FloatTensor, "_hard_negatives embedding_dim"] + ] + random_neg_embeddings: dict[ + CondensedNodeType, + Float32[torch.FloatTensor, "_random_negatives embedding_dim"], + ] # Returns scores for a single anchor node @dataclass class BatchScores: - pos_scores: torch.FloatTensor - hard_neg_scores: torch.FloatTensor - random_neg_scores: torch.FloatTensor + pos_scores: Float32[torch.FloatTensor, "1 positives"] + hard_neg_scores: Float32[torch.FloatTensor, "1 hard_negatives"] + random_neg_scores: Float32[torch.FloatTensor, "1 random_negatives"] # Returns combined scores across all anchor nodes with repeated anchor node embeddings for each positive supervision edge @dataclass class BatchCombinedScores: - repeated_candidate_scores: torch.FloatTensor - positive_ids: torch.LongTensor - hard_neg_ids: torch.LongTensor - random_neg_ids: torch.LongTensor - repeated_query_ids: Optional[torch.LongTensor] + repeated_candidate_scores: Float32[torch.FloatTensor, "queries candidates"] + positive_ids: Int64[torch.LongTensor, "positives"] + hard_neg_ids: Int64[torch.LongTensor, "hard_negatives"] + random_neg_ids: Int64[torch.LongTensor, "random_negatives"] + repeated_query_ids: Optional[Int64[torch.LongTensor, "queries"]] num_unique_query_ids: Optional[int] diff --git a/gigl/src/common/utils/eval_metrics.py b/gigl/src/common/utils/eval_metrics.py index 27f537f3d..d7a8247c9 100644 --- a/gigl/src/common/utils/eval_metrics.py +++ b/gigl/src/common/utils/eval_metrics.py @@ -1,9 +1,12 @@ import torch +from jaxtyping import Float, Float32, Int64 def hit_rate_at_k( - pos_scores: torch.FloatTensor, neg_scores: torch.FloatTensor, ks: torch.LongTensor -) -> torch.FloatTensor: + pos_scores: Float[torch.FloatTensor, "*batch positives"], + neg_scores: Float[torch.FloatTensor, "*batch negatives"], + ks: Int64[torch.LongTensor, "requested_ks"], +) -> Float32[torch.FloatTensor, "requested_ks"]: """Computes Hit Rate @ K metrics for various Ks, evaluating 1+ positives against 1+ negatives. Args: @@ -47,8 +50,9 @@ def hit_rate_at_k( def mean_reciprocal_rank( - pos_scores: torch.FloatTensor, neg_scores: torch.FloatTensor -) -> torch.FloatTensor: + pos_scores: Float[torch.FloatTensor, "*batch positives"], + neg_scores: Float[torch.FloatTensor, "*batch negatives"], +) -> Float32[torch.FloatTensor, ""]: """Computes Mean Reciprocal Rank (MRR), evaluating 1+ positives against 1+ negatives. Args: diff --git a/gigl/transforms/graph_transformer.py b/gigl/transforms/graph_transformer.py index e17ed87ff..085d1d389 100644 --- a/gigl/transforms/graph_transformer.py +++ b/gigl/transforms/graph_transformer.py @@ -60,6 +60,7 @@ from typing import Literal, NamedTuple, Optional, TypedDict import torch +from jaxtyping import Bool, Float, Int64 from torch import Tensor from torch_geometric.data import Data, HeteroData from torch_geometric.typing import NodeType @@ -71,10 +72,10 @@ class SequenceAuxiliaryData(TypedDict): - anchor_bias: Optional[Tensor] - pairwise_bias: Optional[Tensor] - pairwise_relation_indices: Optional[Tensor] - pairwise_nonmissing_indices: Optional[Tensor] + anchor_bias: Optional[Float[Tensor, "anchors sequence"]] + pairwise_bias: Optional[Float[Tensor, "anchors sequence sequence"]] + pairwise_relation_indices: Optional[Int64[Tensor, "relation_edges 4"]] + pairwise_nonmissing_indices: Optional[Int64[Tensor, "nonmissing_edges 3"]] token_input: Optional[TokenInputData] @@ -96,7 +97,7 @@ def heterodata_to_graph_transformer_input( batch_size: int, max_seq_len: int, anchor_node_type: NodeType, - anchor_node_ids: Optional[Tensor] = None, + anchor_node_ids: Optional[Int64[Tensor, "anchors"]] = None, hop_distance: int = 2, sequence_construction_method: Literal["khop", "ppr"] = "khop", include_anchor_first: bool = True, @@ -106,7 +107,11 @@ def heterodata_to_graph_transformer_input( pairwise_attention_bias_attr_names: Optional[list[str]] = None, relation_edge_types: Optional[list[GiGLEdgeType]] = None, sampling_direction: Literal["in", "out"] = "out", -) -> tuple[Tensor, Tensor, SequenceAuxiliaryData]: +) -> tuple[ + Float[Tensor, "anchors sequence feature_dim"], + Bool[Tensor, "anchors sequence"], + SequenceAuxiliaryData, +]: """ Transform a HeteroData object to Graph Transformer sequence input. diff --git a/gigl/utils/sampling.py b/gigl/utils/sampling.py index 5d0ed6a44..e9dad7b76 100644 --- a/gigl/utils/sampling.py +++ b/gigl/utils/sampling.py @@ -3,6 +3,7 @@ from typing import Any, Optional, Union import torch +from jaxtyping import Int64 from gigl.common.logger import Logger from gigl.src.common.types.graph_data import EdgeType, NodeType @@ -136,7 +137,13 @@ class ABLPInputNodes: ) """ - anchor_nodes: torch.Tensor + anchor_nodes: Int64[torch.Tensor, "anchors"] anchor_node_type: NodeType - labels: dict[EdgeType, tuple[torch.Tensor, Optional[torch.Tensor]]] + labels: dict[ + EdgeType, + tuple[ + Int64[torch.Tensor, "anchors _positive_labels_per_anchor"], + Optional[Int64[torch.Tensor, "anchors _negative_labels_per_anchor"]], + ], + ] diff --git a/pyproject.toml b/pyproject.toml index be07400e4..8c17ab51c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ # See https://stackoverflow.com/questions/69759351/error-jupyter-client-kernelspec-nosuchkernel-no-such-kernel-named-python3-occu "ipykernel", "ipython", + "jaxtyping", "kfp>=2.0.0", "matplotlib", "mmh3", @@ -149,6 +150,7 @@ typing-stubs = [ ] test = [ "parameterized==0.9.0", + "typeguard>=2,<3", # Jaxtyping documents known incompatibilities with Typeguard 3 and 4. ] lint = [ "ruff==0.15.10", diff --git a/tests/e2e_tests/e2e_test.py b/tests/e2e_tests/e2e_test.py index 4fa81b070..c51a41283 100644 --- a/tests/e2e_tests/e2e_test.py +++ b/tests/e2e_tests/e2e_test.py @@ -49,6 +49,10 @@ from dataclasses import dataclass, field from typing import Optional +from tests.test_assets.runtime_type_checking import install_runtime_typechecking + +install_runtime_typechecking() + from google.cloud.aiplatform import PipelineJob from gigl import __version__ diff --git a/tests/integration/common/data/export_test.py b/tests/integration/common/data/export_test.py index f83b3fa9e..dc20c68ed 100644 --- a/tests/integration/common/data/export_test.py +++ b/tests/integration/common/data/export_test.py @@ -63,7 +63,7 @@ def test_embedding_export(self, _, should_run_async: bool): with EmbeddingExporter(export_dir=self.embedding_output_dir) as exporter: for i in torch.arange(num_nodes): exporter.add_embedding( - torch.tensor([i]), torch.ones(128, 1) * i, "node" + torch.tensor([i]), torch.ones(1, 128) * i, "node" ) # We also want nested directories to be picked up. @@ -76,7 +76,7 @@ def test_embedding_export(self, _, should_run_async: bool): ) as exporter: for i in torch.arange(num_nodes, num_nodes * 2): exporter.add_embedding( - torch.tensor([i]), torch.ones(128, 1) * i, "node" + torch.tensor([i]), torch.ones(1, 128) * i, "node" ) bq_client = BqUtils() bq_export_table_path = bq_client.join_path( diff --git a/tests/integration/main.py b/tests/integration/main.py index fd1765afd..f72d337b2 100644 --- a/tests/integration/main.py +++ b/tests/integration/main.py @@ -1,5 +1,9 @@ import sys +from tests.test_assets.runtime_type_checking import install_runtime_typechecking + +install_runtime_typechecking() + import gigl.src.common.constants.local_fs as local_fs_constants from gigl.common import LocalUri from gigl.common.utils.test_utils import parse_args, run_tests diff --git a/tests/test_assets/runtime_type_checking.py b/tests/test_assets/runtime_type_checking.py new file mode 100644 index 000000000..f5faa2f0d --- /dev/null +++ b/tests/test_assets/runtime_type_checking.py @@ -0,0 +1,124 @@ +"""Install test-only runtime checks for GiGL tensor Shape Contracts. + +Test launchers call ``install_runtime_typechecking()`` before test discovery. +Jaxtyping then instruments GiGL and example modules as they are imported, and +Typeguard checks only annotations containing Jaxtyping array types when those +calls run. Production imports do not install this hook. +""" + +import atexit +import os +from functools import reduce +from operator import or_ +from types import FunctionType, UnionType +from typing import Any, Final, Optional, Union, cast, get_args, get_origin + +from jaxtyping import AbstractArray, install_import_hook +from typeguard import typechecked + +from gigl.common.logger import Logger + +_SHAPE_CONTRACT_PACKAGES: Final[tuple[str, ...]] = ("gigl", "examples") + +_import_hook: Optional[object] = None +_instrumented_functions: set[str] = set() +logger = Logger() + + +def shape_contract_typechecker(function: FunctionType) -> FunctionType: + """Apply Typeguard only to annotations containing Jaxtyping arrays. + + Args: + function: Function imported from a Shape Contract module. + + Returns: + Function wrapped to enforce only its Shape Contract annotations. + """ + + def contains_shape_contract(annotation: object) -> bool: + if isinstance(annotation, type) and issubclass(annotation, AbstractArray): + return True + return any(contains_shape_contract(arg) for arg in get_args(annotation)) + + def retain_shape_contract(annotation: object) -> object: + if isinstance(annotation, type) and issubclass(annotation, AbstractArray): + return annotation + origin = get_origin(annotation) + args = get_args(annotation) + if not args or origin is None: + return annotation + if origin in (Union, UnionType): + retained_args = tuple( + retain_shape_contract(arg) if contains_shape_contract(arg) else arg + for arg in args + ) + else: + # Existing non-shape members are outside this test-only contract and + # may contain forward references Typeguard cannot resolve here. + retained_args = tuple( + retain_shape_contract(arg) if contains_shape_contract(arg) else Any + for arg in args + ) + if hasattr(annotation, "copy_with"): + return cast(Any, annotation).copy_with(retained_args) + if origin is UnionType: + return reduce(or_, retained_args) + return origin[retained_args] + + annotations = function.__annotations__ + shape_annotations = { + name: retain_shape_contract(annotation) + for name, annotation in annotations.items() + if contains_shape_contract(annotation) + } + if not shape_annotations: + return function + # Typeguard reads annotations again at call time. A private function copy + # keeps its shape-only view from replacing the public API annotations. + checking_function = FunctionType( + function.__code__, + function.__globals__, + function.__name__, + function.__defaults__, + function.__closure__, + ) + checking_function.__annotations__ = shape_annotations + checking_function.__kwdefaults__ = function.__kwdefaults__ + wrapped_function = typechecked(checking_function) + wrapped_function.__annotations__ = annotations + _instrumented_functions.add(f"{function.__module__}.{function.__qualname__}") + return wrapped_function + + +def install_runtime_typechecking() -> None: + """Enable test-only runtime checks for tensor Shape Contracts. + + Jaxtyping's simpler, general-purpose setup would be:: + + install_import_hook( + modules=("gigl", "examples"), + typechecker="typeguard.typechecked", + ) + + That setup makes every annotation in those packages a runtime contract. + GiGL has static-only annotations, such as TypeVars bounded by Protocols that + cannot be used with ``isinstance``. The custom typechecker filters those out + so Typeguard enforces Shape Contracts only. + + Repeated calls are safe. Runtime checking remains scoped to the current + process and modules imported after this function runs. A contract violation + raises ``jaxtyping.TypeCheckError`` at the call site, so an uncaught + violation fails the active test command. + """ + global _import_hook + if _import_hook is None: + _import_hook = install_import_hook( + modules=_SHAPE_CONTRACT_PACKAGES, + typechecker="tests.test_assets.runtime_type_checking.shape_contract_typechecker", + ) + atexit.register( + lambda: logger.info( + f"Shape checks: pid={os.getpid()} count={len(_instrumented_functions)}" + ) + ) + logger.info(f"Shape checks enabled: pid={os.getpid()}") diff --git a/tests/unit/common/data/dataloaders_test.py b/tests/unit/common/data/dataloaders_test.py index d6ccf01f7..62ac6cdab 100644 --- a/tests/unit/common/data/dataloaders_test.py +++ b/tests/unit/common/data/dataloaders_test.py @@ -375,7 +375,7 @@ def test_build_dataset_for_uris(self): "just_node", feature_keys=[], feature_dim=0, - expected_node_ids=torch.empty(0), + expected_node_ids=torch.empty(0, dtype=torch.int64), expected_features=None, expected_label_tensor=None, entity_key="node_id", @@ -384,7 +384,7 @@ def test_build_dataset_for_uris(self): "node_with_features", feature_keys=["foo_feature"], feature_dim=1, - expected_node_ids=torch.empty(0), + expected_node_ids=torch.empty(0, dtype=torch.int64), expected_features=torch.empty(0, 1), expected_label_tensor=None, entity_key="node_id", @@ -393,7 +393,7 @@ def test_build_dataset_for_uris(self): "just_edge", feature_keys=[], feature_dim=0, - expected_node_ids=torch.empty(2, 0), + expected_node_ids=torch.empty(2, 0, dtype=torch.int64), expected_features=None, expected_label_tensor=None, entity_key=("src_node_id", "dst_node_id"), @@ -402,7 +402,7 @@ def test_build_dataset_for_uris(self): "edge_with_features", feature_keys=["foo_feature", "bar_feature"], feature_dim=3, - expected_node_ids=torch.empty(2, 0), + expected_node_ids=torch.empty(2, 0, dtype=torch.int64), expected_features=torch.empty(0, 3), expected_label_tensor=None, entity_key=("src_node_id", "dst_node_id"), @@ -411,7 +411,7 @@ def test_build_dataset_for_uris(self): "node_with_label_only", feature_keys=[], feature_dim=0, - expected_node_ids=torch.empty(0), + expected_node_ids=torch.empty(0, dtype=torch.int64), expected_features=None, expected_label_tensor=torch.empty(0, 1), # 1 label entity_key="node_id", @@ -421,7 +421,7 @@ def test_build_dataset_for_uris(self): "node_with_features_and_label", feature_keys=["foo_feature"], feature_dim=1, - expected_node_ids=torch.empty(0), + expected_node_ids=torch.empty(0, dtype=torch.int64), expected_features=torch.empty(0, 1), # 1 feature expected_label_tensor=torch.empty(0, 1), # 1 label entity_key="node_id", diff --git a/tests/unit/common/utils/feature_quantization/numpy_ops_test.py b/tests/unit/common/utils/feature_quantization/numpy_ops_test.py index a6e6f4136..792f639ed 100644 --- a/tests/unit/common/utils/feature_quantization/numpy_ops_test.py +++ b/tests/unit/common/utils/feature_quantization/numpy_ops_test.py @@ -1,4 +1,5 @@ import numpy as np +from jaxtyping import TypeCheckError from gigl.common.utils.feature_quantization.numpy_ops import quantize_ndarray from tests.test_assets.test_case import TestCase @@ -81,7 +82,7 @@ def test_quantize_ndarray_rejects_invalid_bit_width(self) -> None: quantize_ndarray(np.zeros((1, 1)), bits=3, clip_min=0.0, clip_max=1.0) def test_quantize_ndarray_rejects_non_2d_features(self) -> None: - with self.assertRaises(ValueError): + with self.assertRaises(TypeCheckError): quantize_ndarray(np.zeros((1, 1, 1)), bits=2, clip_min=0.0, clip_max=1.0) def test_quantize_ndarray_requires_multi_bit_clip_bounds(self) -> None: diff --git a/tests/unit/distributed/sampler_test.py b/tests/unit/distributed/sampler_test.py index 2e62dafff..18ebaef10 100644 --- a/tests/unit/distributed/sampler_test.py +++ b/tests/unit/distributed/sampler_test.py @@ -32,12 +32,12 @@ def _build_sampler_input( """Builds a simple ABLPNodeSamplerInput for testing with two edge types.""" node = torch.arange(num_nodes) positive_label_by_edge_types = { - _USER_BUYS_ITEM: torch.arange(100, 100 + num_nodes), - _USER_CLICKS_ITEM: torch.arange(200, 200 + num_nodes), + _USER_BUYS_ITEM: torch.arange(100, 100 + num_nodes).unsqueeze(1), + _USER_CLICKS_ITEM: torch.arange(200, 200 + num_nodes).unsqueeze(1), } negative_label_by_edge_types = { - _USER_BUYS_ITEM: torch.arange(300, 300 + num_nodes), - _USER_CLICKS_ITEM: torch.arange(400, 400 + num_nodes), + _USER_BUYS_ITEM: torch.arange(300, 300 + num_nodes).unsqueeze(1), + _USER_CLICKS_ITEM: torch.arange(400, 400 + num_nodes).unsqueeze(1), } return ABLPNodeSamplerInput( node=node, @@ -50,8 +50,8 @@ def _build_sampler_input( class TestABLPNodeSamplerInput(TestCase): def test_construction_and_properties(self) -> None: node = torch.tensor([10, 20, 30]) - positive_labels = {_USER_BUYS_ITEM: torch.tensor([1, 2, 3])} - negative_labels = {_USER_CLICKS_ITEM: torch.tensor([4, 5, 6])} + positive_labels = {_USER_BUYS_ITEM: torch.tensor([[1], [2], [3]])} + negative_labels = {_USER_CLICKS_ITEM: torch.tensor([[4], [5], [6]])} sampler_input = ABLPNodeSamplerInput( node=node, @@ -94,19 +94,19 @@ def test_getitem_with_tensor_index(self) -> None: self.assertEqual(sliced.input_type, _USER) self.assert_tensor_equality( sliced.positive_label_by_edge_types[_USER_BUYS_ITEM], - torch.tensor([100, 102]), + torch.tensor([[100], [102]]), ) self.assert_tensor_equality( sliced.positive_label_by_edge_types[_USER_CLICKS_ITEM], - torch.tensor([200, 202]), + torch.tensor([[200], [202]]), ) self.assert_tensor_equality( sliced.negative_label_by_edge_types[_USER_BUYS_ITEM], - torch.tensor([300, 302]), + torch.tensor([[300], [302]]), ) self.assert_tensor_equality( sliced.negative_label_by_edge_types[_USER_CLICKS_ITEM], - torch.tensor([400, 402]), + torch.tensor([[400], [402]]), ) def test_getitem_with_list_index(self) -> None: @@ -116,10 +116,11 @@ def test_getitem_with_list_index(self) -> None: self.assertIsInstance(sliced, ABLPNodeSamplerInput) self.assertTrue(torch.equal(sliced.node, torch.tensor([1]))) self.assert_tensor_equality( - sliced.positive_label_by_edge_types[_USER_BUYS_ITEM], torch.tensor([101]) + sliced.positive_label_by_edge_types[_USER_BUYS_ITEM], torch.tensor([[101]]) ) self.assert_tensor_equality( - sliced.negative_label_by_edge_types[_USER_CLICKS_ITEM], torch.tensor([401]) + sliced.negative_label_by_edge_types[_USER_CLICKS_ITEM], + torch.tensor([[401]]), ) def test_share_memory(self) -> None: @@ -165,8 +166,12 @@ def test_prepare_ablp_inputs_dedupes_same_type_seeds_and_keeps_anchors_first( self, ) -> None: sampler = _build_sampler_stub(edge_dir="out") - positive_labels = {_USER_FRIEND_USER: torch.tensor([11, 12, -1, 13])} - negative_labels = {_USER_FRIEND_USER: torch.tensor([13, 14, 10, -1])} + positive_labels = { + _USER_FRIEND_USER: torch.tensor([[11, 12], [-1, 13], [-1, -1]]) + } + negative_labels = { + _USER_FRIEND_USER: torch.tensor([[13, 14], [10, -1], [-1, -1]]) + } sampler_input = ABLPNodeSamplerInput( node=torch.tensor([10, 11, 10]), input_type=_USER, @@ -206,10 +211,10 @@ def test_prepare_ablp_inputs_dedupes_cross_type_supervision_nodes(self) -> None: node=torch.tensor([4, 5]), input_type=_USER, positive_label_by_edge_types={ - _USER_BUYS_ITEM: torch.tensor([20, 21, 20, -1]) + _USER_BUYS_ITEM: torch.tensor([[20, 21], [20, -1]]) }, negative_label_by_edge_types={ - _USER_BUYS_ITEM: torch.tensor([21, 22, -1, 20]) + _USER_BUYS_ITEM: torch.tensor([[21, 22], [-1, 20]]) }, ) diff --git a/tests/unit/main.py b/tests/unit/main.py index 83b3b75d4..cd016814e 100644 --- a/tests/unit/main.py +++ b/tests/unit/main.py @@ -1,5 +1,9 @@ import sys +from tests.test_assets.runtime_type_checking import install_runtime_typechecking + +install_runtime_typechecking() + import gigl.src.common.constants.local_fs as local_fs_constants from gigl.common import LocalUri from gigl.common.utils.test_utils import parse_args, run_tests diff --git a/tests/unit/nn/models_test.py b/tests/unit/nn/models_test.py index 0d62bf104..11728b6f9 100644 --- a/tests/unit/nn/models_test.py +++ b/tests/unit/nn/models_test.py @@ -14,7 +14,10 @@ from gigl.nn.models import LightGCN, LinkPredictionGNN from gigl.src.common.types.graph_data import NodeType from gigl.types.graph import DEFAULT_HOMOGENEOUS_NODE_TYPE -from tests.test_assets.distributed.utils import get_process_group_init_method +from tests.test_assets.distributed.utils import ( + create_test_process_group, + get_process_group_init_method, +) from tests.test_assets.test_case import TestCase # Embedding table name for default homogeneous node type @@ -42,10 +45,10 @@ def forward( "Output node types must be specified for heterogeneous data" ) return { - node_type: torch.tensor([1.0, 2.0]) for node_type in output_node_types + node_type: torch.tensor([[1.0, 2.0]]) for node_type in output_node_types } else: - return torch.tensor([1.0, 2.0]) + return torch.tensor([[1.0, 2.0]]) class DummyDecoder(nn.Module): @@ -59,7 +62,7 @@ def __init__(self): def forward( self, query_embeddings: torch.Tensor, candidate_embeddings: torch.Tensor ) -> torch.Tensor: - return query_embeddings + candidate_embeddings + return torch.mm(query_embeddings, candidate_embeddings.T) class TestLinkPredictionGNN(TestCase): @@ -73,7 +76,7 @@ def test_forward_homogeneous(self): data = Data() result = model.forward(data, self.device) assert isinstance(result, torch.Tensor) - self.assert_tensor_equality(result, torch.tensor([1.0, 2.0])) + self.assert_tensor_equality(result, torch.tensor([[1.0, 2.0]])) def test_forward_heterogeneous_with_node_types(self): encoder = DummyEncoder() @@ -85,7 +88,7 @@ def test_forward_heterogeneous_with_node_types(self): assert isinstance(result, dict) self.assertEqual(set(result.keys()), set(output_node_types)) for node_type in output_node_types: - self.assert_tensor_equality(result[node_type], torch.tensor([1.0, 2.0])) # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. + self.assert_tensor_equality(result[node_type], torch.tensor([[1.0, 2.0]])) # ty: ignore[invalid-argument-type] TODO(ty-torch-keyed-access): fix ty false positives for torch-backed keyed container access. def test_forward_heterogeneous_missing_node_types(self): encoder = DummyEncoder() @@ -99,10 +102,10 @@ def test_decode(self): encoder = DummyEncoder() decoder = DummyDecoder() model = LinkPredictionGNN(encoder, decoder) - q = torch.tensor([1.0, 2.0]) - c = torch.tensor([3.0, 4.0]) + q = torch.tensor([[1.0, 2.0]]) + c = torch.tensor([[3.0, 4.0]]) result = model.decode(q, c) - self.assert_tensor_equality(result, torch.tensor([4.0, 6.0])) + self.assert_tensor_equality(result, torch.tensor([[11.0]])) def test_encoder_property(self): encoder = DummyEncoder() @@ -117,9 +120,7 @@ def test_decoder_property(self): self.assertIs(model.decoder, decoder) def test_for_ddp(self): - torch.distributed.init_process_group( - rank=0, world_size=1, init_method=get_process_group_init_method() - ) + create_test_process_group() self.addCleanup(torch.distributed.destroy_process_group) encoder = DummyEncoder() decoder = DummyDecoder() @@ -132,9 +133,7 @@ def test_for_ddp(self): self.assertTrue(hasattr(ddp_model.decoder, "module")) def test_unwrap_from_ddp(self): - torch.distributed.init_process_group( - rank=0, world_size=1, init_method=get_process_group_init_method() - ) + create_test_process_group() self.addCleanup(torch.distributed.destroy_process_group) encoder = DummyEncoder() decoder = DummyDecoder() diff --git a/tests/unit/src/common/models/layers/loss_test.py b/tests/unit/src/common/models/layers/loss_test.py index 4be356ca3..0376667d7 100644 --- a/tests/unit/src/common/models/layers/loss_test.py +++ b/tests/unit/src/common/models/layers/loss_test.py @@ -167,9 +167,9 @@ def test_loss_value(self): ) ) - def test_empty_loss(self): + def test_empty_loss(self) -> None: loss_fn = RetrievalLoss(remove_accidental_hits=True) - query_ids = torch.empty(0) + query_ids = torch.empty(0, dtype=torch.long) empty_scores = torch.empty((0, 5)) candidate_ids = torch.tensor([0, 1, 2, 3, 4]) expected_loss = torch.tensor(0.0) diff --git a/tests/unit/src/common/types/task_inputs_test.py b/tests/unit/src/common/types/task_inputs_test.py new file mode 100644 index 000000000..51c9276e9 --- /dev/null +++ b/tests/unit/src/common/types/task_inputs_test.py @@ -0,0 +1,40 @@ +import torch + +from gigl.src.common.types.graph_data import CondensedEdgeType, CondensedNodeType +from gigl.src.common.types.task_inputs import ( + BatchCombinedScores, + BatchEmbeddings, + BatchScores, +) +from tests.test_assets.test_case import TestCase + + +class TaskInputsTest(TestCase): + def test_empty_task_inputs_retain_matrix_rank(self) -> None: + condensed_edge_type = CondensedEdgeType(0) + condensed_node_type = CondensedNodeType(0) + + batch_embeddings = BatchEmbeddings( + query_embeddings=torch.FloatTensor(0, 8), + repeated_query_embeddings={condensed_edge_type: torch.FloatTensor(0, 8)}, + pos_embeddings={condensed_edge_type: torch.FloatTensor(0, 8)}, + hard_neg_embeddings={condensed_edge_type: torch.FloatTensor(0, 8)}, + random_neg_embeddings={condensed_node_type: torch.FloatTensor(0, 8)}, + ) + batch_scores = BatchScores( + pos_scores=torch.FloatTensor(1, 0), + hard_neg_scores=torch.FloatTensor(1, 0), + random_neg_scores=torch.FloatTensor(1, 0), + ) + combined_scores = BatchCombinedScores( + repeated_candidate_scores=torch.FloatTensor(0, 0), + positive_ids=torch.LongTensor([]), + hard_neg_ids=torch.LongTensor([]), + random_neg_ids=torch.LongTensor([]), + repeated_query_ids=torch.LongTensor([]), + num_unique_query_ids=0, + ) + + self.assertEqual(batch_embeddings.query_embeddings.shape, (0, 8)) + self.assertEqual(batch_scores.pos_scores.shape, (1, 0)) + self.assertEqual(combined_scores.repeated_candidate_scores.shape, (0, 0)) diff --git a/uv.lock b/uv.lock index 98e483d80..3ee49cc02 100644 --- a/uv.lock +++ b/uv.lock @@ -715,6 +715,7 @@ dependencies = [ { name = "google-cloud-storage" }, { name = "ipykernel" }, { name = "ipython" }, + { name = "jaxtyping" }, { name = "kfp" }, { name = "matplotlib" }, { name = "mmh3" }, @@ -795,6 +796,7 @@ dev = [ { name = "sphinx-rtd-theme" }, { name = "sphinx-tabs" }, { name = "ty" }, + { name = "typeguard" }, { name = "types-psutil" }, { name = "types-pyyaml" }, { name = "types-requests" }, @@ -830,6 +832,7 @@ lint = [ ] test = [ { name = "parameterized" }, + { name = "typeguard" }, ] typing-stubs = [ { name = "pandas-stubs" }, @@ -855,6 +858,7 @@ requires-dist = [ { name = "hydra-core", marker = "extra == 'experimental'", specifier = "==1.3.2" }, { name = "ipykernel" }, { name = "ipython" }, + { name = "jaxtyping" }, { name = "kfp", specifier = ">=2.0.0" }, { name = "matplotlib" }, { name = "mmh3" }, @@ -921,6 +925,7 @@ dev = [ { name = "sphinx-rtd-theme", specifier = "==2.0.0" }, { name = "sphinx-tabs", specifier = "==3.4.5" }, { name = "ty", specifier = "==0.0.31" }, + { name = "typeguard", specifier = ">=2,<3" }, { name = "types-psutil", specifier = "==7.0.0.20250401" }, { name = "types-pyyaml", specifier = "~=6.0.12" }, { name = "types-requests", specifier = "==2.31.0.6" }, @@ -954,7 +959,10 @@ lint = [ { name = "ruff", specifier = "==0.15.10" }, { name = "ty", specifier = "==0.0.31" }, ] -test = [{ name = "parameterized", specifier = "==0.9.0" }] +test = [ + { name = "parameterized", specifier = "==0.9.0" }, + { name = "typeguard", specifier = ">=2,<3" }, +] typing-stubs = [ { name = "pandas-stubs", specifier = "==2.2.2.240807" }, { name = "types-psutil", specifier = "==7.0.0.20250401" }, @@ -1755,6 +1763,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d9/33/1f075bf72b0b747cb3288d011319aaf64083cf2efef8354174e3ed4540e2/ipython_pygments_lexers-1.1.1-py3-none-any.whl", hash = "sha256:a9462224a505ade19a605f71f8fa63c2048833ce50abc86768a0d81d876dc81c", size = 8074, upload-time = "2025-01-17T11:24:33.271Z" }, ] +[[package]] +name = "jaxtyping" +version = "0.3.11" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wadler-lindig" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/91/c1/091b8852bd7cbf50bd655543c8506033cf4029300c67f8c176c1286879a9/jaxtyping-0.3.11.tar.gz", hash = "sha256:b09c14acf6686feb9e0df5b0d8c6e7c5b6f8d36bf059ee54cd522a186c2ef050", size = 46489, upload-time = "2026-06-13T18:35:23.167Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/38/c66bbdc5047f4776c2bd3e47e5295a350e3fa44d5b8942105e71c2a876a0/jaxtyping-0.3.11-py3-none-any.whl", hash = "sha256:8a4bedc4e3f963fa82df41bd13c7ebc2bad925601eb48614c65798f21329d4e3", size = 56593, upload-time = "2026-06-13T18:35:22.01Z" }, +] + [[package]] name = "jedi" version = "0.19.2" @@ -4452,6 +4472,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/88/39/bca669095ccf0a400af941fdf741578d4c2d6719f1b7f10e6dbec10aa862/ty-0.0.31-py3-none-win_arm64.whl", hash = "sha256:e9cb15fad26545c6a608f40f227af3a5513cb376998ca6feddd47ca7d93ffafa", size = 10590392, upload-time = "2026-04-15T15:47:57.968Z" }, ] +[[package]] +name = "typeguard" +version = "2.13.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3a/38/c61bfcf62a7b572b5e9363a802ff92559cb427ee963048e1442e3aef7490/typeguard-2.13.3.tar.gz", hash = "sha256:00edaa8da3a133674796cf5ea87d9f4b4c367d77476e185e80251cc13dfbb8c4", size = 40604, upload-time = "2021-12-10T21:09:39.158Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/bb/d43e5c75054e53efce310e79d63df0ac3f25e34c926be5dffb7d283fb2a8/typeguard-2.13.3-py3-none-any.whl", hash = "sha256:5e3e3be01e887e7eafae5af63d1f36c849aaa94e3a0112097312aabfa16284f1", size = 17605, upload-time = "2021-12-10T21:09:37.844Z" }, +] + [[package]] name = "types-protobuf" version = "6.32.1.20250918" @@ -4608,6 +4637,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/79/0c/c05523fa3181fdf0c9c52a6ba91a23fbf3246cc095f26f6516f9c60e6771/virtualenv-20.35.4-py3-none-any.whl", hash = "sha256:c21c9cede36c9753eeade68ba7d523529f228a403463376cf821eaae2b650f1b", size = 6005095, upload-time = "2025-10-29T06:57:37.598Z" }, ] +[[package]] +name = "wadler-lindig" +version = "0.1.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1e/67/cbae4bf7683a64755c2c1778c418fea96d00e34395bb91743f08bd951571/wadler_lindig-0.1.7.tar.gz", hash = "sha256:81d14d3fe77d441acf3ebd7f4aefac20c74128bf460e84b512806dccf7b2cd55", size = 15842, upload-time = "2025-06-18T07:00:42.843Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/96/04e7b441807b26b794da5b11e59ed7f83b2cf8af202bd7eba8ad2fa6046e/wadler_lindig-0.1.7-py3-none-any.whl", hash = "sha256:e3ec83835570fd0a9509f969162aeb9c65618f998b1f42918cfc8d45122fe953", size = 20516, upload-time = "2025-06-18T07:00:41.684Z" }, +] + [[package]] name = "wcwidth" version = "0.2.14"