Enable scaled matmul tests on ROCm #678
magaonka-amd
wants to merge
6 commits into
ROCm:rocm-jaxlib-v0.8.0
from
magaonka-amd:enable-scaled-matmul-tests-v0.8.0
Add missing sharding test case from upstream
3f4a69e9
Fix overlapping partition specs in scaled_matmul partitioner
5ab92579
Skip cuDNN and Blackwell checks on ROCm devices
fe3ccbd7
Add platform-specific custom call target constants
d498963a
Add ROCm specific HLO calls and fallback handling
9e49d3f4
Enable scaled_matmul tests on ROCm platform
816d5f4e
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Login to write a write a comment.
Login via GitHub