Repository: google/jax

Issue:
Default lowering pass for lu_p has accuracy issue.
### Description

The background is, I am developing another device backend for JAX, I found the default lowering pass for lu_p has accuracy issue. It can be reproduced on NV platform.
By default if runnning `python tests/linalg_test.py NumpyLinalgTest.testDetGrad1` it works well.
But if I comment cuda lowering pass [here](https://github.com/google/jax/blob/jaxlib-v0.4.20/jax/_src/lax/linalg.py#L1387-L1390), it goes default pass [here](https://github.com/google/jax/blob/jaxlib-v0.4.20/jax/_src/lax/linalg.py#L1378), then running UT raises error:
```
AssertionError:
Not equal to tolerance rtol=0.1, atol=0.1
JVP primal
x and y nan location mismatch:
 x: array(nan, dtype=float32)
 y: array(0., dtype=float32)
```

### What jax/jaxlib version are you using?

jax 0.4.20, jaxlib 0.4.20+cuda11.cudnn86

### Which accelerator(s) are you using?

GPU

### Additional system info?

python 3.9

### NVIDIA GPU info

_No response_


Discussion:
Hi - thanks for the report. It looks like this `_lu_python` implementation is dead code: it is never invoked, because there are platform-specific lowering rules for `cpu`, `cuda`, `rocm`, and `tpu`, which are all the platforms JAX supports.

With that in mind, I think we could probably delete this code, as it is unused. Would that resolve your issue?
Ah, actually it is not dead code, because it is used in the `jax2tf` lowering path here: https://github.com/google/jax/blob/42ae8432185bf03f61ddd2e7bc279d3abb5247fd/jax/experimental/jax2tf/jax2tf.py#L3299
The issue is that the `_lu_python` implementation does not handle the gradient of the determinant for zero-valued inputs. Here's a self-contained reproduction of the issue (run on CPU):
```python
import jax.numpy as jnp
from jax._src.interpreters import mlir
from jax.lax import linalg
import jax.test_util as jtu

x = jnp.zeros((3, 3))
assert x.device().platform == 'cpu'
del mlir._platform_specific_lowerings['cpu'][linalg.lu_p]  # this ensures that the default lowering is used.
jtu.check_grads(jnp.linalg.det, (x,), order=1)
# Not equal to tolerance rtol=0.002, atol=0.002
# JVP primal
# x and y nan location mismatch:
#  x: array(nan, dtype=float32)
#  y: array(0., dtype=float32)
```
Even more minimal:
```python
import jax.numpy as jnp
from jax._src.lax import linalg
x = jnp.zeros((2, 2))

print(linalg.lu(x)[0])
# [[0. 0.]
#  [0. 0.]]

print(linalg._lu_python(x)[0])
# [[ 0.  0.]
#  [nan nan]]
```
The root cause is the division operation here: https://github.com/google/jax/blob/b07d176419552dd7147da8002aa80fe4b42bc917/jax/_src/lax/linalg.py#L1166-L1168

When `a` is all zeros, then this is `0 / 0` which becomes `NaN`

Resolving patch:
diff --git a/jax/_src/lax/linalg.py b/jax/_src/lax/linalg.py
--- a/jax/_src/lax/linalg.py
+++ b/jax/_src/lax/linalg.py
@@ -1159,13 +1159,13 @@ def body(k, state):
     else:
       magnitude = ufuncs.abs(a[:, k])
     i = jnp.argmax(jnp.where(m_idx >= k, magnitude, -jnp.inf))
-    pivot = pivot.at[k].set(i)
+    pivot = pivot.at[k].set(i.astype(pivot.dtype))
     a = a.at[[k, i],].set(a[[i, k],])
     perm = perm.at[[i, k],].set(perm[[k, i],])
 
     # a[k+1:, k] /= a[k, k], adapted for loop-invariant shapes
     x = a[k, k]
-    a = a.at[:, k].set(jnp.where(m_idx > k, a[:, k] / x, a[:, k]))
+    a = a.at[:, k].set(jnp.where((m_idx > k) & (x != 0), a[:, k] / x, a[:, k]))
 
     # a[k+1:, k+1:] -= jnp.outer(a[k+1:, k], a[k, k+1:])
     a = a - jnp.where((m_idx[:, None] > k) & (n_idx[None, :] > k),


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -23,7 +23,7 @@
 import scipy.linalg
 import scipy as osp
 
-from absl.testing import absltest
+from absl.testing import absltest, parameterized
 
 import jax
 from jax import jit, grad, jvp, vmap
@@ -32,6 +32,7 @@
 from jax import scipy as jsp
 from jax._src.lib import version as jaxlib_version
 from jax._src import config
+from jax._src.lax import linalg as lax_linalg
 from jax._src import test_util as jtu
 from jax._src import xla_bridge
 from jax._src.numpy.util import promote_dtypes_inexact
@@ -109,7 +110,6 @@ def testDetGrad(self, shape, dtype):
     jtu.check_grads(jnp.linalg.det, (a,), 2, atol=1e-1, rtol=1e-1)
     # make sure there are no NaNs when a matrix is zero
     if len(shape) == 2:
-      pass
       jtu.check_grads(
         jnp.linalg.det, (jnp.zeros_like(a),), 1, atol=1e-1, rtol=1e-1)
     else:
@@ -1279,6 +1279,13 @@ def testLuOfSingularMatrix(self):
     p, l, u = jsp.linalg.lu(x)
     self.assertAllClose(x, np.matmul(p, np.matmul(l, u)))
 
+  @parameterized.parameters(lax_linalg.lu, lax_linalg._lu_python)
+  def testLuOnZeroMatrix(self, lu):
+    # Regression test for https://github.com/google/jax/issues/19076
+    x = jnp.zeros((2, 2), dtype=np.float32)
+    x_lu, _, _ = lu(x)
+    self.assertArraysEqual(x_lu, x)
+
   @jtu.sample_product(
     shape=[(1, 1), (4, 5), (10, 5), (10, 10), (6, 7, 7)],
     dtype=float_types + complex_types,