mirror of
https://github.com/ROCm/jax.git
synced 2025-04-16 03:46:06 +00:00

This also requires some changes on the XLA side for the paths and such to have any effect. The plugin init now looks for a `rocm` python package install and extracts ROCm toolkit paths from the python packages that it finds. We hand these into the XLA portions of the plugin via environment variables to avoid changing any interfaces like protobuf or PJRT C APIs. We also have to patch the rpath in the shared object files included in the plugin and kernel wheels so they look relative to their install path just like the cuda based plugin objects do. Some other changes are fixing missing dynamic link libraries and also adding an optional feature target to pull in rocm python dependencies for the plugin.