Merge pull request #6144 from google:issue6121

PiperOrigin-RevId: 364018304

jax/_src/numpy/linalg.py

@@ -202,7 +202,7 @@ def _cofactor_solve(a, b):
            "a=[..., m, m] and b=[..., m, m]; got a={} and b={}")
     raise ValueError(msg.format(a_shape, b_shape))
   if a_shape[-1] == 1:
-    return a[0, 0], b
+    return a[..., 0, 0], b
   # lu contains u in the upper triangular matrix and l in the strict lower
   # triangular matrix.
   # The diagonal of l is set to ones without loss of generality.

tests/linalg_test.py

@@ -119,6 +119,12 @@ def testDetGrad(self, shape, dtype):
       a[0] = 0
       jtu.check_grads(jnp.linalg.det, (a,), 1, atol=1e-1, rtol=1e-1)
 
+  def testDetGradIssue6121(self):
+    f = lambda x: jnp.linalg.det(x).sum()
+    x = jnp.ones((16, 1, 1))
+    jax.grad(f)(x)
+    jtu.check_grads(f, (x,), 2, atol=1e-1, rtol=1e-1)
+
   def testDetGradOfSingularMatrixCorank1(self):
     # Rank 2 matrix with nonzero gradient
     a = jnp.array([[ 50, -30,  45],