Add test-time tensor shape contracts - #748
Conversation
There was a problem hiding this comment.
These text fixtures were flagged as having invalid shapes w.r.t. to the prod path. Updates here pass shape checks.
| enable_residual_topup: bool = True, | ||
| num_neighbors_per_hop: int = 100_000, | ||
| degree_tensors: Union[torch.Tensor, dict[NodeType, torch.Tensor]], | ||
| degree_tensors: Union[ |
There was a problem hiding this comment.
An example of where the union type is actually encoding two different shapes without a requirement that they match: the homogenous type (first type) and heterogeneous (second type) need not be equal.
There was a problem hiding this comment.
In this case, the shape of all nodes in the dict shall not be equal, e.x. {"a": torch.ones(2), "b": torch.ones(5)} is valid an expected.
What's the jaxtyping behavior here?
There was a problem hiding this comment.
As I understand it, typing the example like dict[str, Float[Tensor, "nodes"]] would force all values in the dictionary to have matching shape.
We specifically use _nodes (note the underscore _) which effectively skips any shape checks for the type. This allows us to communicate the expected semantics without actually enforcing a shape check (since they need not match).
| self, query_embeddings: torch.Tensor, candidate_embeddings: torch.Tensor | ||
| ) -> torch.Tensor: | ||
| return query_embeddings + candidate_embeddings | ||
| return torch.mm(query_embeddings, candidate_embeddings.T) |
There was a problem hiding this comment.
This mirrors the production decoder: [queries, embedding_dim] @ [embedding_dim, candidates] returns pairwise [queries, candidates] scores. Addition returns embeddings, not decoder scores.
kmontemayor2-sc
left a comment
kmontemayor2-sc
left a comment
There was a problem hiding this comment.
Neat! Thanks for exploring Jacob :) I left some comments :)
I guess for this we'd need to be careful that we only enable runtime shape checking for tests?
| enable_residual_topup: bool = True, | ||
| num_neighbors_per_hop: int = 100_000, | ||
| degree_tensors: Union[torch.Tensor, dict[NodeType, torch.Tensor]], | ||
| degree_tensors: Union[ |
There was a problem hiding this comment.
In this case, the shape of all nodes in the dict shall not be equal, e.x. {"a": torch.ones(2), "b": torch.ones(5)} is valid an expected.
What's the jaxtyping behavior here?
| @@ -37,11 +42,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, Int[torch.Tensor, "anchors positive_labels_per_anchor"]]: | |||
| return self._positive_label_by_edge_types | |||
There was a problem hiding this comment.
In this case, is anchors bound per instance of ABLPSamplerInput?
e.x.
in1 = Input(node: torch.ones(10), pos: {e: torch.zeroes(10, 2))
in2 = Input(node: torch.ones(10), pos: {e: torch.zeroes(10, 3))
Can we later distinguish between the label size for different objects?
My question about dict key sizing from above still stands.
There was a problem hiding this comment.
This is a good question. Jaxtyping bindings are scoped to single call so your sample code would pass. This property in particular is essentially just validating a 2d integral output tensor.
Jaxtyping does not do any analysis across function calls and object lifetimes. It only validates between input and output tensors of the same function call.
For something that can check across function calls, I think we'd need to look into static type checks like pyrefly which has an experimental feature for tensor shape checks.
|
|
||
| from tests.test_assets.runtime_type_checking import install_runtime_typechecking | ||
|
|
||
| install_runtime_typechecking() |
There was a problem hiding this comment.
Hmmm, in some of our tests we fork / spawn new processes, would we need to call this again in those?
There was a problem hiding this comment.
I'd have to double check! We specifically install the runtime type check in test module main, so the spawned process should end up re-running it automaitcally since its at module scope
| from beartype import beartype | ||
| from jaxtyping import AbstractArray, install_import_hook | ||
|
|
||
| _SHAPE_CONTRACT_MODULES: Final[tuple[str, ...]] = ( |
There was a problem hiding this comment.
Hmmm, so we'd need to add this for all source modules we want to test?
Do you think there's a way to enable this for all gigl/ and examples/ and then see if the calls seem as expected?
E.g. add some new make check_tensor?
There was a problem hiding this comment.
We should be able to directly put package prefixes to make this more maintainable.
It's important to realize that this is a runtime type checker, so only the code that runs is checked. So having make check_tensor is nice, but in practice we need something to actually run (e.g. the tests) in order for it to do anything.
2 participants
Purpose of this PR
This PR adds targeted, runtime-checkable tensor contracts for: loader and sampler inputs, public model
forwardanddecodemethods, loss interfaces, and task-result containers.Jaxtyping lets an annotation declare tensor dtype, rank, fixed dimensions, and relationships between named dimensions. For example,
Float[Tensor, "queries embedding_dim"]andFloat[Tensor, "candidates embedding_dim"]require matchingembedding_dim; decoder outputFloat[Tensor, "queries candidates"]then documents both output axes. This makes malformed tensors fail close to the boundary.The contracts are intentionally test-only. Unit, integration, and E2E launchers install a Jaxtyping hook before test discovery. Listed modules imported afterwards are instrumented; arguments are checked before execution and returns afterwards. An uncaught violation raises
jaxtyping.TypeCheckError, which fails the test command. Production execution does not enable this mechanism, and this PR does not expose it as a user API.Why this is useful:
[anchors, labels_per_anchor], not flat vectorsPotential downsides:
Note: