jax
5d1c0bc9 - [ROCm] Enable HLO module transform registration for GPU backends

Commit
52 days ago
[ROCm] Enable HLO module transform registration for GPU backends register_hlo_module_transformation(platforms=None) was silently skipping GPU plugin clients because the default platforms list only contained factory keys (e.g. "rocm") while PJRT plugin clients report client.platform as "gpu". Add "gpu" as an explicit entry so the backend initialization hook fires for all GPU plugin clients. Also enable the xla_transform_test for the GPU backend.
Author
Parents
Loading