rocm_jax/jax_plugins
Mathew Odden f965454982 Add support for ROCm wheel based install (#283)
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.
2025-03-27 14:32:34 -05:00
..