PR #20288: cuda_root_path: Find cuda libraries when installed with co…

…nda packages

Imported from GitHub PR #20288

This fix emerged when looking in solving jax-ml/jax#24604 . In a nutshell, the official cuda package for conda (both in the `conda-forge` and `nvidia` conda channels) install the CUDA libraries in a different location with respect to PyPI packages, so the logic to find them needs to be augmented to be able to find the CUDA libraries when installed from conda packages.

I did not tested this with a tensorflow build, but probably this will also help in solving tensorflow/tensorflow#56927 .

xref: conda-forge/tensorflow-feedstock#408
xref: conda-forge/jaxlib-feedstock#288
Copybara import of the project:

--
a2ce85c by Silvio Traversaro <[email protected]>:

cuda_root_path: Find cuda libraries when installed with conda packages

Merging this change closes #20288

FUTURE_COPYBARA_INTEGRATE_REVIEW=#20288 from traversaro:fixloadcudaconda a2ce85c
PiperOrigin-RevId: 717411600

xla/tsl/platform/default/cuda_root_path.cc

@@ -76,6 +76,17 @@ std::vector<std::string> CandidateCudaRoots() {
     // Also add the path to the copy of libdevice.10.bc that we include within
     // the Python wheel.
     roots.emplace_back(io::JoinPath(dir, "cuda"));
+
+    // In case cuda was installed with nvidia's official conda packages, we also
+    // include the root prefix of the environment in the candidate roots dir,
+    // we assume that the lib binaries are either in the python package's root
+    // dir or in a 'python' subdirectory, as done by the previous for. python
+    // packages on non-Windows platforms are installed in
+    // $CONDA_PREFIX/lib/python3.12/site-packages/pkg_name, so if we want
+    // to add $CONDA_PREFIX to the candidate roots dirs we need to add
+    // ../../../..
+    for (auto path : {"../../../..", "../../../../.."})
+      roots.emplace_back(io::JoinPath(dir, path));
   }
 #endif  // defined(PLATFORM_POSIX) && !defined(__APPLE__)