From 88f8050e4261e05a64431c04edce01fbc0cb079c Mon Sep 17 00:00:00 2001 From: Philipp Holl Date: Fri, 10 Jan 2025 16:07:46 +0100 Subject: [PATCH] Fix _linear_solve_forward() --- phiml/math/_optimize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/phiml/math/_optimize.py b/phiml/math/_optimize.py index ff3b6e3..2bbe4fe 100644 --- a/phiml/math/_optimize.py +++ b/phiml/math/_optimize.py @@ -714,7 +714,7 @@ def _linear_solve_forward(y: Tensor, matrix_offset = reshaped_tensor(approx_matrix_vals, [batch_dims]) matrix_offset = math.where(solve.rank_deficiency > 0, matrix_offset, 0).native([batch_dims]) method = solve.method - if not callable(native_lin_op) and is_sparse(native_lin_op) and y.default_backend.name == 'torch' and preconditioner and not all_available(y): + if not callable(native_lin_op) and is_sparse(native_lin_op) and y_tensor.default_backend.name == 'torch' and preconditioner and not all_available(y): warnings.warn(f"Preconditioners are not supported for sparse {method} in {y.default_backend} JIT mode. Disabling preconditioner. Use Jax or TensorFlow to enable preconditioners in JIT mode.", RuntimeWarning) preconditioner = None if not callable(native_lin_op) and is_sparse(native_lin_op) and not all_available(y) and not method.startswith('scipy-') and isinstance(preconditioner, IncompleteLU):