If LAPACK kernels are linked statically, don't override them from scipy.

PiperOrigin-RevId: 839343548

Author: hawkinsp

jax/BUILD

@@ -245,7 +245,7 @@ py_library_providing_imports_info(
         "//jax/_src:xla_metadata",
         "//jax/_src:xla_metadata_lib",
         "//jax/_src/lib",
-    ] + py_deps("numpy") + py_deps("scipy") + py_deps("opt_einsum") + py_deps("flatbuffers") + jax_extra_deps,
+    ] + py_deps("numpy") + py_deps("opt_einsum") + py_deps("flatbuffers") + jax_extra_deps,
 )
 
 pytype_strict_library(

jaxlib/cpu/lapack.cc

@@ -27,6 +27,9 @@ using ::xla::ffi::DataType;
 
 void GetLapackKernelsFromScipy() {
   static absl::once_flag initialized;
+  if (lapack_kernels_initialized) {
+    return;
+  }
   // For reasons I'm not entirely sure of, if the import_ call is done inside
   // the call_once scope, we sometimes observe deadlocks in the test suite.
   // However it probably doesn't do much harm to just import them a second time,

jaxlib/cpu/lapack_kernels.cc

@@ -48,6 +48,8 @@ XLA_FFI_REGISTER_ENUM_ATTR_DECODING(jax::schur::Sort);
 
 namespace jax {
 
+bool lapack_kernels_initialized = false;
+
 template <typename T>
 inline T CastNoOverflow(int64_t value, std::string_view source = __FILE__) {
   auto result = MaybeCastNoOverflow<T>(value, source);

jaxlib/cpu/lapack_kernels.h

@@ -30,6 +30,8 @@ limitations under the License.
 
 namespace jax {
 
+extern bool lapack_kernels_initialized;
+
 struct MatrixParams {
   enum class Side : char { kLeft = 'L', kRight = 'R' };
   enum class UpLo : char { kLower = 'L', kUpper = 'U' };

jaxlib/cpu/lapack_kernels_using_lapack.cc

@@ -173,6 +173,7 @@ static auto init = []() -> int {
   AssignKernelFn<TridiagonalSolver<ffi::DataType::C64>>(cgtsv_);
   AssignKernelFn<TridiagonalSolver<ffi::DataType::C128>>(zgtsv_);
 
+  lapack_kernels_initialized = true;
   return 0;
 }();
 

jaxlib/gpu/hybrid.cc

@@ -29,6 +29,9 @@ namespace nb = nanobind;
 
 void GetLapackKernelsFromScipy() {
   static absl::once_flag initialized;
+  if (lapack_kernels_initialized) {
+    return;
+  }
   // For reasons I'm not entirely sure of, if the import_ call is done inside
   // the call_once scope, we sometimes observe deadlocks in the test suite.
   // However it probably doesn't do much harm to just import them a second time,