jax
75d1d883 - Fix plugin init script

Commit
15 days ago
Fix plugin init script Add ROCm __init__.py to plugin wheel sources Include jax_plugins/rocm/__init__.py in the wheel_sources static_srcs so it can be copied during the wheel build process. This file contains critical ROCm initialization logic including: - ROCm path configuration (bitcode, linker) - AMD GPU detection via KFD - Shared memory size validation for multi-GPU setups Fixes build error where build_gpu_kernels_wheel.py couldn't find the __init__.py file to copy into the plugin wheel.
Author
Committer
Parents
Loading