-
Notifications
You must be signed in to change notification settings - Fork 487
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Don't set $TPU_LIBRARY_PATH during import #5698
Conversation
I am a bit confuse after this pr, what will happen if user still set |
We'll still take Internally, we can use |
|
||
We load libtpu.so in the following order of precedence: | ||
|
||
1. User-set $TPU_LIBRARY_PATH | ||
2. libtpu.so included in torch_xla/lib | ||
3. libtpu-nightly pip package | ||
|
||
Sets $PTXLA_TPU_LIBRARY_PATH if path is inferred by us to prevent conflicts | ||
with other frameworks. This env var will be removed in a future version. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I wonder under what condition can we remove this env var.
* Don't set $TPU_LIBRARY_PATH during import * remove chekck for new env var so people don't use it
This PR regress the profiler. I cannot take multi-host profilers after this change. I'm going to revert it. |
This reverts commit 146f2a0.
* Don't set $TPU_LIBRARY_PATH during import * remove chekck for new env var so people don't use it
* Don't set $TPU_LIBRARY_PATH during import * remove chekck for new env var so people don't use it
…torch#5731) This reverts commit 146f2a0.
* Don't set $TPU_LIBRARY_PATH during import * remove chekck for new env var so people don't use it
…torch#5731) This reverts commit 146f2a0.
* Don't set $TPU_LIBRARY_PATH during import * remove chekck for new env var so people don't use it
* Don't set $TPU_LIBRARY_PATH during import * remove chekck for new env var so people don't use it
JAX and PyTorch/XLA are both overriding
$TPU_LIBRARY_PATH
during import, leading to confusing issues that depend on import order, e.g. #5625get_library_path
function to get that path from the libtpu package.$TPU_LIBRARY_PATH
as an override for compatibilityCorresponding fix in JAX: jax-ml/jax@b81a3e1
The issue will be resolved when both JAX and PyTorch/XLA release new packages that include these two changes.
$PTXLA_TPU_LIBRARY_PATH
is a temporary hack. I expect we'll have a better way to get the PJRT Plugin path in the near future.cc @jyingl3