Add NaN check to prevent false singular matrix errors

src/ATen/native/xpu/mkl/BatchLinearAlgebra.cpp

@@ -19,8 +19,10 @@
 #include <ATen/native/xpu/mkl/BatchLinearAlgebra.h>
 #include <ATen/ops/_linalg_check_errors.h>
 #include <ATen/ops/_linalg_check_errors_native.h>
+#include <ATen/ops/arange.h>
 #include <ATen/ops/empty.h>
 #include <ATen/ops/from_blob.h>
+#include <ATen/ops/isnan.h>
 #include <ATen/ops/zeros_like.h>
 
 #include <comm/SYCLContext.h>
@@ -526,6 +528,15 @@ void lu_solve_mkl(
     const Tensor& pivots,
     const Tensor& B,
     TransposeType trans) {
+  // NaN check: if LU or B contains NaN, fill B with NaN and return
+  if (at::isnan(LU).any().item<bool>() || at::isnan(B).any().item<bool>()) {
+    AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(
+        B.scalar_type(), "lu_solve_mkl_nan_fill", [&] {
+          B.fill_(std::numeric_limits<scalar_t>::quiet_NaN());
+        });
+    return;
+  }
+
   AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(LU.scalar_type(), "lu_solve_xpu", [&] {
     apply_lu_solve_xpu_<scalar_t>(LU, pivots, B, trans);
   });
@@ -545,6 +556,22 @@ void lu_factor_mkl(
       pivot,
       "linalg.lu_factor: LU without pivoting is not implemented on the XPU");
 
+  // NaN check: if input contains NaN, return NaN output with valid pivots
+  if (at::isnan(LU).any().item<bool>()) {
+    info.zero_();
+    // Fill pivots with default sequence [1, 2, 3, ..., min(m, n)]
+    int64_t min_mn = std::min(LU.size(-2), LU.size(-1));
+    auto default_pivots = at::arange(
+        1, min_mn + 1, pivots.options().dtype(at::kInt));
+    pivots.copy_(default_pivots.expand_as(pivots));
+    // Fill LU with NaN
+    AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(
+        LU.scalar_type(), "lu_factor_mkl_nan_fill", [&] {
+          LU.fill_(std::numeric_limits<scalar_t>::quiet_NaN());
+        });
+    return;
+  }
+
   // handle the info
   Tensor info_ = at::zeros_like(info, Device(at::kCPU));
   int32_t* info_data = info_.data_ptr<int32_t>();