Repository: google/jax
Issue: Gradient of jnp.linalg.det fails

Hi,

I am quite new to JAX but there seems to be an issue with taking the derivative through `jnp.linalg.det`. The following code fails (it works for functions like `jnp.linalg.slogdet`):
```python
import jax
import jax.numpy as jnp
key = jax.random.PRNGKey(42)
jax.grad(lambda x: jnp.linalg.det(x).sum())(jax.random.normal(key, (16, 1, 1)))
```
```
---------------------------------------------------------------------------
FilteredStackTrace                        Traceback (most recent call last)
<ipython-input-64-74aa1dd68db7> in <module>
      1 key = jax.random.PRNGKey(42)
----> 2 jax.grad(lambda x: jnp.linalg.det(x).sum())(jax.random.normal(key, (16, 1, 1)))

<ipython-input-64-74aa1dd68db7> in <lambda>(x)
      1 key = jax.random.PRNGKey(42)
----> 2 jax.grad(lambda x: jnp.linalg.det(x).sum())(jax.random.normal(key, (16, 1, 1)))

FilteredStackTrace: TypeError: Custom JVP rule must produce primal and tangent outputs with equal shapes and dtypes, but got float32[1] and float32[16] respectively.

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)
<ipython-input-64-74aa1dd68db7> in <module>
      1 key = jax.random.PRNGKey(42)
----> 2 jax.grad(lambda x: jnp.linalg.det(x).sum())(jax.random.normal(key, (16, 1, 1)))

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    137   def reraise_with_filtered_traceback(*args, **kwargs):
    138     try:
--> 139       return fun(*args, **kwargs)
    140     except Exception as e:
    141       if not is_under_reraiser(e):

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/api.py in grad_f(*args, **kwargs)
    758   @api_boundary
    759   def grad_f(*args, **kwargs):
--> 760     _, g = value_and_grad_f(*args, **kwargs)
    761     return g
    762 

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    137   def reraise_with_filtered_traceback(*args, **kwargs):
    138     try:
--> 139       return fun(*args, **kwargs)
    140     except Exception as e:
    141       if not is_under_reraiser(e):

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    821     tree_map(partial(_check_input_dtype_grad, holomorphic, allow_int), dyn_args)
    822     if not has_aux:
--> 823       ans, vjp_py = _vjp(f_partial, *dyn_args)
    824     else:
    825       ans, vjp_py, aux = _vjp(f_partial, *dyn_args, has_aux=True)

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/api.py in _vjp(fun, has_aux, *primals)
   1894   if not has_aux:
   1895     flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-> 1896     out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
   1897     out_tree = out_tree()
   1898   else:

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux)
    112 def vjp(traceable, primals, has_aux=False):
    113   if not has_aux:
--> 114     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    115   else:
    116     out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
     99   _, in_tree = tree_flatten(((primals, primals), {}))
    100   jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
--> 101   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
    102   out_primals_pvals, out_tangents_pvals = tree_unflatten(out_tree(), out_pvals)
    103   assert all(out_primal_pval.is_known() for out_primal_pval in out_primals_pvals)

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate)
    504   with core.new_main(JaxprTrace) as main:
    505     fun = trace_to_subjaxpr(fun, main, instantiate)
--> 506     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    507     assert not env
    508     del main, fun, env

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    164 
    165     try:
--> 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:
    168       # Some transformations yield from inside context managers, so we have to

<ipython-input-64-74aa1dd68db7> in <lambda>(x)
      1 key = jax.random.PRNGKey(42)
----> 2 jax.grad(lambda x: jnp.linalg.det(x).sum())(jax.random.normal(key, (16, 1, 1)))

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/custom_derivatives.py in __call__(self, *args, **kwargs)
    215     flat_jvp, out_tree2 = _flatten_jvp(jvp, in_tree)
    216     if config.omnistaging_enabled:
--> 217       out_flat = custom_jvp_call_p.bind(flat_fun, flat_jvp, *args_flat)
    218       _, out_tree = lu.merge_linear_aux(out_tree1, out_tree2)
    219     else:

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/custom_derivatives.py in bind(self, fun, jvp, *args)
    281     tracers = map(top_trace.full_raise, args)  # type: ignore
    282     with core.maybe_new_sublevel(top_trace):
--> 283       outs = top_trace.process_custom_jvp_call(self, fun, jvp, tracers)  # type: ignore
    284     _, env_trace_todo = lu.merge_linear_aux(env_trace_todo1, env_trace_todo2)
    285     return _apply_todos(env_trace_todo, map(core.full_lower, outs))

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/interpreters/ad.py in process_custom_jvp_call(self, _, __, f_jvp, tracers)
    350     # currently handle float0s
    351     tangents_in = map(replace_float0s, primals_in, tangents_in)
--> 352     outs = f_jvp.call_wrapped(*it.chain(primals_in, tangents_in))
    353     primals_out, tangents_out = split_list(outs, [len(outs) // 2])
    354     tangents_out = map(recast_to_float0, primals_out, tangents_out)

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    177     while stack:
    178       gen, out_store = stack.pop()
--> 179       ans = gen.send(ans)
    180       if out_store is not None:
    181         ans, side = ans

~/miniconda3/envs/jax/lib/python3.8/site-packages/jax/custom_derivatives.py in _flatten_jvp(in_tree, *args)
    259       msg = ("Custom JVP rule must produce primal and tangent outputs with "
    260              "equal shapes and dtypes, but got {} and {} respectively.")
--> 261       raise TypeError(msg.format(av1.str_short(), av2.str_short()))
    262     else:
    263       msg = ("Custom JVP rule must produce primal and tangent outputs with "

TypeError: Custom JVP rule must produce primal and tangent outputs with equal shapes and dtypes, but got float32[1] and float32[16] respectively.
```

Changed files:
tests/linalg_test.py, jax/_src/numpy/linalg.py

diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
index e4ec6be3c..7df4825ff 100644
--- a/jax/_src/numpy/linalg.py
+++ b/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.
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index bf016c92f..eda711228 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -119,6 +119,12 @@ class NumpyLinalgTest(jtu.JaxTestCase):
       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],