Add test-time tensor shape contracts

[jchmura-sc]

Purpose of this PR

This PR adds targeted, runtime-checkable tensor contracts for: loader and sampler inputs, public model forward and decode methods, 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"] and Float[Tensor, "candidates embedding_dim"] require matching embedding_dim; decoder output Float[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:

• documents the tensor contracts agents and users must satisfy
• checks dtype, rank, fixed axes, and repeated named-axis equality in exercised code paths
• caught the previously undocumented ABLP label contract: labels are padded [anchors, labels_per_anchor], not flat vectors

Potential downsides:

• Developers must know a boundary’s rank and axis meaning. Unknown numeric sizes are fine: named axes bind runtime values. If rank or semantics are unclear, do not guess a contract.
• Incorrect or overly strict contracts create false test failures and add maintenance during API changes.

Note:

• low-level message-passing operations remain uncontracted because they do not expose one stable tensor shape
• dictionary-valued contracts validate tensor values, but cannot express equality by matching dictionary key (e.g. happens a lot for heterogeneous graph data structures)

[jchmura-sc]

These text fixtures were flagged as having invalid shapes w.r.t. to the prod path. Updates here pass shape checks.

[kmontemayor2-sc]

Hmmm, I wonder if this is "wrong"? The code still works as expected right? I'm not sure if it's bad form to treat [x, 1] and [x] shaped tensors as equivalent, but maybe we could? WDYT?

[kmontemayor2-sc]

We could use #num_labels here right, to indicate the size can be 1? https://docs.kidger.site/jaxtyping/api/advanced-features/#jaxtyping.AbstractArray

[jchmura-sc]

I'm not sure if it's bad form to treat [x, 1] and [x] shaped tensors as equivalent, but maybe we could? WDYT?

I would say that they are broadly equivalent in that they are always broadcast-able to each other, but they could have different behaviors in some cases:

• torch.cat(..., dim=1), or some other op where we specify dim=1 would break under [x]
• torch.softmax(..., dim=-1) would produce different normalization for [x, 1] or [x]
• torch.stack(...) would produce 2d or 3d tensor depending on dimension of input.

From a docs perspective it could also help inform whether you need a .squeeze() or not. I've had situations where I had missed that, the code ran fine, but the computed loss was completely wrong.

[jchmura-sc]

We could use #num_labels here right, to indicate the size can be 1? https://docs.kidger.site/jaxtyping/api/advanced-features/#jaxtyping.AbstractArray

Where exactly are you suggesting this? An annotation like [foo, #num_labels] allow shape [x, y] or [x, 1] but not [x].

Also, if a function only has a single typed argument, adding the # wouldnt' change anything since num_labels (no #) can bind to any size.

[kmontemayor2-sc]

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?

[kmontemayor2-sc]

I wonder if we can extend AbstractArray so that we can also tie "anchors" into this class? https://docs.kidger.site/jaxtyping/api/advanced-features/#jaxtyping.AbstractArray

e.g. verify that this shape and the shape of input_seeds are compatible?

[jchmura-sc]

We can actually do something slightly simpler, using symbol expressions "{...}"

Something like

input_seeds: Int64[torch.Tensor, "{inputs.node.shape[0]}"],

seems to do what you suggested? This gets resolved and will raise an exception if they are incompatible.

Of course, this only runs when jaxtyping hooks is installed - i.e. testing.

[kmontemayor2-sc]

Hmmm, I wonder if this is "wrong"? The code still works as expected right? I'm not sure if it's bad form to treat [x, 1] and [x] shaped tensors as equivalent, but maybe we could? WDYT?

[kmontemayor2-sc]

We could use #num_labels here right, to indicate the size can be 1? https://docs.kidger.site/jaxtyping/api/advanced-features/#jaxtyping.AbstractArray

[jchmura-sc]

/unit_test

[github-actions]

GiGL Automation

@ 20:21:42UTC : 🔄 Scala Unit Test started.

@ 20:32:44UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 20:21:43UTC : 🔄 C++ Unit Test started.

@ 20:23:48UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 20:21:45UTC : 🔄 Python Unit Test started.

@ 21:44:34UTC : ✅ Workflow completed successfully.

[jchmura-sc]

/unit_test

[github-actions]

GiGL Automation

@ 22:26:09UTC : 🔄 Python Unit Test started.

@ 23:57:49UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 22:26:10UTC : 🔄 Scala Unit Test started.

@ 22:37:49UTC : ✅ Workflow completed successfully.

[github-actions]

GiGL Automation

@ 22:26:12UTC : 🔄 C++ Unit Test started.

@ 22:28:11UTC : ✅ Workflow completed successfully.