Nx.Testing: vectorized tensor support + diff diagnostics on failure

[blasphemetheus]

Summary

Nx.Testing.assert_equal/2 and Nx.Testing.assert_all_close/3 now handle vectorized tensors cleanly and produce a numeric diagnostic (max absolute / max relative difference) when an assertion fails. Motivated by a pain point in the vectorized grad test suite that landed in #1731: when two tensors differ by a tiny amount (e.g. 1 ULP) that both inspect to the same string, the old failure message showed the same tensor text twice with no way to see what actually differed.

What changes

• tensor_equal?/2 explicitly guards vec-axes and shape mismatches before running the element-wise comparison. Two tensors with different vectorized axes or shapes always return false instead of raising or relying on broadcast semantics.
• assert_all_close/3 gains the same vec-axes/shape guards at the top so mismatched inputs produce a structural error message instead of an opaque Nx.all_close failure.
• Both helpers now call a new diagnose_difference/2 helper in the flunk block. It:
  • Describes structural mismatches directly (vectorized_axes differ: ..., shapes differ: ...).
  • For tensors that DO share structure, computes max absolute and max relative difference across all elements (including vec axes, via devectorize-before-reduce). Output looks like:

    max absolute difference: 0.5
    max relative difference: 0.2
    

• assert_all_close/3 failure messages also include the atol and rtol values that were being checked against (previously hidden).
• Defensive fallback: if the diff computation itself raises (mixed complex/real types, unusual promotions), diagnose_difference/2 catches and returns an empty string so the baseline inspect output is still shown — no failure mode where the diagnostic itself becomes an obstacle.

Before / after

Before — a 1-ULP disagreement that both tensors inspect the same:

Tensor assertion failed.
left: #Nx.Tensor<f32[3] [1.0, 2.0, 3.0]>
right: #Nx.Tensor<f32[3] [1.0, 2.0, 3.0]>

(You have to attach a debugger or IO.inspect the raw binaries to figure out what went wrong.)

After — the same failure:

Tensor assertion failed.

left:

#Nx.Tensor<f32[3] [1.0, 2.0, 3.0]>

right:

#Nx.Tensor<f32[3] [1.0, 2.0, 3.0]>

max absolute difference: 1.1920929e-7
max relative difference: 5.9604645e-8

Visually, the tensors still match, but now you can see they differ by exactly one f32 ULP.

Tests

New nx/test/nx/testing_test.exs — 12 tests covering:

• Both helpers pass on bit-identical tensors and within-tolerance tensors (plain and vectorized).
• Both helpers' error messages include the max absolute difference on a genuine disagreement.
• Both helpers' error messages include a clear diagnostic when vec axes differ.
• assert_all_close error message includes max-diff for vectorized tensors with close-but-not-equal values (the 1-ULP case).
• assert_equal handles NaN-equality correctly (pre-existing behavior, covered for regression).

All 12 pass locally.

Scope / non-goals

• Out of scope: changes to assert_equal's NaN semantics, changes to Nx.all_close itself, or introducing new public helpers. This is strictly an improvement to the failure-reporting behavior of the two existing assertion helpers.
• Co-authorship: credit Paulo Valente for review-time suggestions during PR Support vectorized gradients via boundary devectorization (#1533) #1731 that shaped the vec-axes guarding approach (he caught that the assertion helper was the source of a confusing test failure rather than the unit under test).

Related

• Support vectorized gradients via boundary devectorization (#1533) #1731 — vectorized grad boundary approach. Caused the confusing 1-ULP vs inspect-identical failures that motivated this cleanup.

[blasphemetheus]

Update on CI after bundling the sharding fix:

• exla 1.17.3: green ✓ (was red on sharding)
• exla 1.18.4: red on a formatting issue — Elixir 1.18's formatter wants blank lines between my new assert_equals and the trailing comments on sharding_test.exs. Fixed in 8c3686ea, should go green on next run.
• torchx (both versions): red on two tests in torchx/test/torchx/nx_test.exs (output_permutation and input_dilation) that turn out to have the same class of latent shape bug as the sharding cases — the test asserts result.shape == {1, 3, 3, 2} on one line, then passes a value literal of shape {3, 3, 2} to assert_all_close on the next, and the old broadcast-tolerant assertion hid the contradiction.

Bundled the torchx fix here too (commit 163f7da0) on the same "get CI green" rationale as the sharding bundling. Same minimal shape: wrap each expected literal in the missing leading [...] brackets to match the shape already explicitly asserted via result.shape ==. Happy to split into a separate PR if you'd rather keep #1733's scope tight to nx/

[blasphemetheus]

remaining failing ci is just precision issue