Make _cuda_path more reliable
#24139
Conversation
hartikainen
force-pushed
the
fix-cuda_path
branch
from
b5e1736
to
c09bb9f
Compare
|
Ah, you're right! I got distracted by the comment about |
jax/_src/lib/__init__.py
Outdated
| cuda_nvcc_path = pathlib.Path(cuda_nvcc.__file__).parent | ||
| return str(cuda_nvcc_path) | ||
|
|
||
| if (path := _try_jaxlib_relative_path()) is not None: |
There was a problem hiding this comment.
I think I'd be comfortable simply removing the jaxlib-relative path. It never was the right thing to do, given, for example, you might have installed cuda_nvcc in the system package root but jaxlib in a different package root.
There was a problem hiding this comment.
Sounds good! I'll remove it. Do you think it would make sense to add a check for $CUDA_ROOT? That should fix Filippo's issue in #22590. I'm happy to add that here while we're at it.
There was a problem hiding this comment.
Yes, I think that would be a great idea.
There was a problem hiding this comment.
I just pushed the $CUDA_ROOT change. Let me know if anything needs to be tweaked.
hartikainen
force-pushed
the
fix-cuda_path
branch
from
6fbc715
to
7e2d601
Compare
There was a problem hiding this comment.
Two more changes:
- I think you need to add
nvidiato the mypy unknown module list here:Line 13 in e8cea0d
- Please squash your commits.
Thanks!
- Remove jax-relative module path test - Use `$CUDA_ROOT` environment variable if available - Use `cuda_nvcc` module's path if installed
hartikainen
force-pushed
the
fix-cuda_path
branch
from
7e2d601
to
1ea8e3c
Compare
|
Thanks for the comments, @hawkinsp. I just squashed my changes and rebased the branch. |
See my comment in #22590 (comment). This defaults to the original behavior of finding
cuda_nvccpath as relative to jaxlib, and only applies the new check if the original way fails. Pulling the path directly fromcuda_nvccseems a bit less error prone and, to me at least, feels preferable over the original way. But keeping this order should be more backwards compatible.