Fix gemm/gemv fusion under cast_policy=numpy+floatX and make tests policy-aware#2276
Open
maresb wants to merge 5 commits into
Open
Fix gemm/gemv fusion under cast_policy=numpy+floatX and make tests policy-aware#2276maresb wants to merge 5 commits into
maresb wants to merge 5 commits into
Conversation
Under cast_policy="numpy+floatX" the +/-1.0 scale seeded by _gemm_canonicalize autocasts to float64, fails the upcast check in _gemm_from_factored_list, and the whole gemm/gemv candidate is silently discarded (e.g. v2 + dot(v1, m) keeps a separate Add instead of folding beta into Gemv). Build literal scales with the matrix dtype instead. No-op under the default custom policy; test_gemv2 is parametrized on cast_policy to cover both. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5
Add a _autocast_int_dtype() helper and use it for the TestExpLog/
TestSqrSqrt switch-constant expectations, drop the now-obsolete
("int8", 0, 1) case from the bitwise-or test, and pin the float32
literals in the two sigmoid rewrite tests. Green under both policies.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5
Pin the +2.0 and /2 fusion constants to np.float32 so the fused output
stays float32, and assert the honest per-policy out_dtype for the
multi-output cases 72/74 via the {"custom": ..., "numpy+floatX": ...}
convention.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5
maresb
force-pushed
the
cast-policy-numpy-floatx
branch
from
2b8430b to
7d860d8
Compare
Pin the inner-graph integer literals that would otherwise promote under numpy+floatX (scan 2*x, scalar loop ** 2 / + 1, gradient x * 2, pytorch Elemwise-ScalarLoop x * 2 / x1 * 3), and make the scalar test_constant and numba float-param dtype expectations derive from the active cast_policy. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5
The hessian benchmark's inner W**2 and cyclical reduction's step_num + 1 promote under numpy+floatX (float32 -> float64 and int32 -> int64), so scan rejects the state dtype and IfElse rejects the branches; pin both literals to int8. The autocast_float_as doctest only makes sense under custom, so scope it with change_flags. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5
maresb
force-pushed
the
cast-policy-numpy-floatx
branch
2 times, most recently
from
aa7a1e0 to
3f5795d
Compare
maresb
changed the title
Collaborator
Author
|
Tests all pass with I have popped the |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
I had some expiring Fable credits last night, so I wanted to see if it could fix the cast policy. I honestly can't say I really understand much myself. Feel free to close in case it's garbage. I suspect it didn't run all the tests, we'll see how far the CI gets.
Description
Bugfix + test hardening, so that the test suite now passes if we were to flip cast_policy. No default change here.
The bugfix: under
cast_policy="numpy+floatX",_gemm_from_factored_listautocasts the ±1.0 canonicalization scale to float64, fails the upcast check, and silently discards the gemm/gemv candidate —v2 + dot(v1, m)keeps a separate Add instead of folding β into Gemv. Literal scales now adopt the matrix dtype (same idiom as #2245/#2247). No-op under the defaultcustompolicy;test_gemv2is parametrized oncast_policyand fails without the fix.The rest is tests/docs: policy-aware expectations or pinned literals so every test passes under both policies while still exercising the same behavior; multi-output fusion cases assert the honest per-policy
out_dtype(the #2248 dict convention); theautocast_float_asdoctest scopes itself tocustom, the only policy where it exists.Verified green under both
customandnumpy+floatX:test_math+test_elemwise355 passed each; gradient/scalar/loop/blas 142; scan benchmarks, pytorch ScalarLoop, and numba rv-param tests all pass both ways.Related Issue
Type of change
🤖 Generated with Claude Code
https://claude.ai/code/session_01FWXc59vTjFTLDybuDBypB5