Improve tests assertion

[skyw]

Add assert_equal and assertion for identity, orthogonality, etc.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR centralises test assertion helpers into a new tests/_comparison.py module, then migrates all 16 test files to use them. The net effect is cleaner, more readable tests with no change to assertion semantics for the vast majority of cases.

• assert_equal wraps torch.testing.assert_close(rtol=0, atol=0) via functools.partial, replacing ~40 repetitive inline calls across the suite.
• assert_close_to_identity / assert_close_to_orthogonal replace bespoke inline Gram-matrix construction and multi-call assert_close patterns with a single, better-documented helper that also produces more informative failure messages by separating diagonal and off-diagonal tolerances.
• A new test_polar_express_16steps_almost_orthogonal test is added in test_muon_utils.py to verify semi-orthogonality of the polar-express Newton-Schulz output.

Confidence Score: 4/5

Safe to merge; the migration is purely mechanical and the new helpers are correct. The one item worth a follow-up before merging is the unexplained 10× tolerance widening in test_polargrad.py.

All assert_equal substitutions preserve exact bitwise semantics. The assert_close_to_identity and assert_close_to_orthogonal helpers are logically correct, including the rtol/atol equivalence in the test_spel migration. The only unresolved question is the tolerance loosening in test_polargrad.py (atol/rtol widened from 1e-5 to 1e-4 with no explanation), which could be hiding a real precision regression or just a noisy test that needed stabilising.

tests/test_polargrad.py — the tolerance change at line 103 lacks a comment explaining whether it reflects a known float32 precision floor or was used to silence intermittent failures.

Important Files Changed

Filename	Overview
tests/_comparison.py	New helper module: assert_equal (bitwise), assert_close_to_identity (separate diagonal/off-diagonal atols), and assert_close_to_orthogonal (Gram-matrix-based). Logic is correct and well-documented.
tests/test_polargrad.py	Replaces assert_close(atol=0,rtol=0) with assert_equal; also silently loosens the RightPolarGradOrthFnTest tolerance by 10× (1e-5→1e-4) with no explanation.
tests/test_muon_utils.py	Renames one test, replaces bitwise asserts with assert_equal, and adds a new test_polar_express_16steps_almost_orthogonal that verifies semi-orthogonality via assert_close_to_orthogonal.
tests/test_spel.py	Replaces verbose inline diagonal/off-diagonal identity check with assert_close_to_identity; semantics are preserved (rtol=0.06 vs atol=0.06 are numerically equivalent for expected value 1).
tests/test_soap_utils.py	Replaces repeated torch.eye construction + assert_close identity checks with assert_close_to_identity, and bitwise permutation asserts with assert_equal.
tests/test_scalar_optimizers.py	Mechanical replacement of assert_close(atol=0,rtol=0) with assert_equal throughout; msg lambdas are preserved correctly via functools.partial passthrough.
tests/test_distributed_muon_utils_cpu.py	One-line assert_equal substitution for tensor-parallel Newton-Schulz output comparison.
tests/test_distributed_rekls_cpu.py	Replaces assert_close(atol=0,rtol=0) with assert_equal for tensor-parallel REKLS output comparison.
tests/test_distributed_soap_utils_cpu.py	Replaces three bitwise assert_close calls with assert_equal for gathered grad and Kronecker factor checks.
tests/test_eig_utils.py	Single assert_equal substitution for conjugate-matrix bitwise identity check.
tests/test_normalized_optimizer.py	Two assert_equal substitutions for momentum buffer and accumulated buffer bitwise checks.
tests/test_orthogonalized_optimizer.py	Four assert_equal substitutions across OrthogonalizedOptimizerTest and MuonTest.
tests/test_triton_kernels.py	Four assert_equal substitutions for tsyrk integer-input kernel output comparisons.
tests/test_utils_modules.py	Six assert_equal substitutions for Conv1dFlatWeights forward and backward pass comparisons.
tests/test_weight_decay_mixin.py	Ten assert_equal substitutions covering all WeightDecayMixin test cases.
tests/test_soap.py	Switches eigenbasis orthogonality check to assert_close_to_identity and bitwise multi-stream equality check to assert_equal.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    CMP["tests/_comparison.py"]
    AE["assert_equal\n(functools.partial of assert_close\natol=rtol=0)"]
    ACI["assert_close_to_identity\n(diag vs ones, off-diag vs zeros\nwith separate atols)"]
    ACO["assert_close_to_orthogonal\n(builds Gram matrix,\nthen calls assert_close_to_identity)"]

    CMP --> AE
    CMP --> ACI
    CMP --> ACO

    AE --> T1["test_distributed_muon_utils_cpu\ntest_distributed_rekls_cpu\ntest_distributed_soap_utils_cpu\ntest_eig_utils / test_muon_utils\ntest_normalized_optimizer\ntest_orthogonalized_optimizer\ntest_polargrad / test_scalar_optimizers\ntest_soap / test_soap_utils\ntest_triton_kernels\ntest_utils_modules / test_weight_decay_mixin"]
    ACI --> T2["test_soap\ntest_soap_utils\ntest_spel"]
    ACO --> T3["test_muon_utils\n(new test_polar_express_16steps)"]

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    CMP["tests/_comparison.py"]
    AE["assert_equal\n(functools.partial of assert_close\natol=rtol=0)"]
    ACI["assert_close_to_identity\n(diag vs ones, off-diag vs zeros\nwith separate atols)"]
    ACO["assert_close_to_orthogonal\n(builds Gram matrix,\nthen calls assert_close_to_identity)"]

    CMP --> AE
    CMP --> ACI
    CMP --> ACO

    AE --> T1["test_distributed_muon_utils_cpu\ntest_distributed_rekls_cpu\ntest_distributed_soap_utils_cpu\ntest_eig_utils / test_muon_utils\ntest_normalized_optimizer\ntest_orthogonalized_optimizer\ntest_polargrad / test_scalar_optimizers\ntest_soap / test_soap_utils\ntest_triton_kernels\ntest_utils_modules / test_weight_decay_mixin"]
    ACI --> T2["test_soap\ntest_soap_utils\ntest_spel"]
    ACO --> T3["test_muon_utils\n(new test_polar_express_16steps)"]

Loading

_{Reviews (1): Last reviewed commit: "Add assert function for identity and ort..." | Re-trigger Greptile}

[greptile-apps]

Unexplained tolerance relaxation

The tolerance on RightPolarGradOrthFnTest was silently widened by 10× (from 1e-5 to 1e-4). If this was triggered by a flaky CI failure, it is worth adding a comment documenting whether the looser bound reflects observed float32 error or was simply the minimum needed to make the test pass consistently.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[skyw]

@greptileai it was a flaky test, ignore.

[greptile-apps]

Got it, makes sense — flaky tests from float32 precision are a common reason to relax tolerances. I'll drop this.

Would you like me to remember a rule: "Do not flag tolerance relaxations in tests when they are used to fix flaky CI failures"?

[skyw]

/ok to test 9708732

[github-actions]

Test Results

   81 files  ±0    151 suites  ±0   1m 47s ⏱️ +3s
1 166 tests +2  1 166 ✅ +2  0 💤 ±0  0 ❌ ±0 
2 714 runs  +4  2 714 ✅ +4  0 💤 ±0  0 ❌ ±0 

Results for commit 9708732. ± Comparison against base commit 9ab7b03.

This pull request removes 3 and adds 5 tests. Note that renamed tests count towards both.

__main__.TestNewtonSchulz ‑ test_newtonschulz5_svd_close0 (512, 512)
__main__.TestNewtonSchulz ‑ test_newtonschulz5_svd_close1 (512, 256)
__main__.TestNewtonSchulz ‑ test_newtonschulz5_svd_close2 (256, 512)

__main__.TestNewtonSchulz ‑ test_newtonschulz5_close_to_svd0 (512, 512)
__main__.TestNewtonSchulz ‑ test_newtonschulz5_close_to_svd1 (512, 256)
__main__.TestNewtonSchulz ‑ test_newtonschulz5_close_to_svd2 (256, 512)
__main__.TestNewtonSchulz ‑ test_polar_express_16steps_almost_orthogonal0 (size=(512, 256))
__main__.TestNewtonSchulz ‑ test_polar_express_16steps_almost_orthogonal1 (size=(256, 512))

♻️ This comment has been updated with latest results.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!