jax
e5e7c2c5 - Simplify _get_library_path to only check for xla_rocm_plugin.so

Commit
72 days ago
Simplify _get_library_path to only check for xla_rocm_plugin.so ROCm always uses wheels for testing, and the wheel build renames pjrt_c_api_gpu_plugin.so to xla_rocm_plugin.so, so we don't need the fallback logic for local builds or runfiles.
Author
Parents
Loading