mirror of
https://github.com/ROCm/jax.git
synced 2025-04-19 05:16:06 +00:00
a03d6e6613
_tpu_ext.so dynamically links in libjaxlib_mlir_capi.so (in jaxlib/mlir/_mlir_libs), so needs to include jaxlib/mlir/_mlir_libs in its RPATH or similar on other platforms. We achieve this by moving _tpu_ext.cc to jaxlib/mlir/_mlir_libs so it can use the same linkopts as other mlir targets that depend on libjaxlib_mlir_capi.so. In particular, we want this to work correctly across platforms, and it's not clear if Windows supports RPATH-like functionality beyond the current directory. PiperOrigin-RevId: 551372130