Make empirical NTK support list of primitive outputs of length 1

PiperOrigin-RevId: 502904434
google · Feb 10, 2023 · 5854a14 · 5854a14
1 parent db1c240
commit 5854a14
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/neural_tangents/_src/empirical.py b/neural_tangents/_src/empirical.py
@@ -1991,7 +1991,7 @@ def _backprop_step(
   """Adapted from `jax.interpreters.ad`."""
   invals = map(functools.partial(_read_primal, primal_env), eqn.invars)
   cts_in = map(read_cotangent, eqn.outvars)
-  if not eqn.primitive.multiple_results:
+  if len(cts_in) == 1:
     cts_in = cts_in[0]
   else:
     raise NotImplementedError(