Jaxtyping Shape Contracts - #748
Conversation
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?
|
/unit_test |
GiGL Automation@ 20:21:42UTC : 🔄 @ 20:32:44UTC : ✅ Workflow completed successfully. |
GiGL Automation@ 20:21:43UTC : 🔄 @ 20:23:48UTC : ✅ Workflow completed successfully. |
GiGL Automation@ 20:21:45UTC : 🔄 @ 21:44:34UTC : ✅ Workflow completed successfully. |
…s' into jchmura/jaxtyping-shape-contracts
|
/unit_test |
GiGL Automation@ 17:30:27UTC : 🔄 @ 17:40:16UTC : ✅ Workflow completed successfully. |
GiGL Automation@ 17:30:30UTC : 🔄 |
GiGL Automation@ 17:30:31UTC : 🔄 @ 17:32:47UTC : ✅ Workflow completed successfully. |
kmontemayor2-sc
left a comment
kmontemayor2-sc
left a comment
There was a problem hiding this comment.
Neat! I think I'm fine with these changes:)
BTW, do you know how much more coverage is untyped?
Can we get a review from @yliu2-sc for the GT shapes as well?
jchmura-sc
changed the title
@kmontemayor2-sc When running python unit tests, 71 functions across the codebase are shape checked. An agent found 334 eligible functions inside |
|
/e2e_test |
GiGL Automation@ 20:55:12UTC : 🔄 |
3 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.