Repository: google/jax

Issue:
Spurious Exception: Tracer can't be used with raw numpy functions?
I cannot seem to make this minimum working example any smaller.  It seems like some combination of factors is causing this error, but I have no idea where to start.
```
import jax
import jax.scipy.special as jss
from jax import numpy as jnp


@jax.custom_jvp
def log_normalizer(q):
    return jnp.logaddexp(q[0], 0.0)


def nat_to_exp(q):
    return jss.expit(q)


@log_normalizer.defjvp
def ln_jvp(primals, tangents):
    x, = primals
    x_dot, = tangents
    y = log_normalizer(x)
    y_dot = nat_to_exp(x) * x_dot
    return y, y_dot


def cross_entropy_loss(p, q):
    # Bernoulli cross_entropy(p, q)
    p_dot_q = jnp.sum(p * q)
    ln = log_normalizer(q)
    return -p_dot_q + ln - jnp.zeros(())


# p are expectation parameters of a Bernoulli distribution corresponding to
# probability 0.4.
p = jnp.array([0.4])

# q are natural parameters of a Bernoulli distribution corresponding to
# log-odds 0, which is probability 0.5.
q = jnp.array([0.0])

print(jax.grad(cross_entropy_loss, 1)(p, q))
```


Discussion:
Wow, this is a terrible error message we're raising. Like, so bad.

It's actually a shape error:

```
> /usr/local/google/home/mattjj/packages/jax/issue2502.py(22)ln_jvp()
     21     import ipdb; ipdb.set_trace()
---> 22     return y, y_dot
     23

ipdb> p y
DeviceArray(0.6931472, dtype=float32)
ipdb> p y_dot
Traced<ShapedArray(float32[1]):JaxprTrace(level=0/0)>
```

That is, `y` and `y_dot` are a scalar and a 1D array with length 1, respectively. This isn't checked in the custom_jvp machinery from #2026, and instead it falls through a strange cascade of steps ultimately leading to us calling `onp.shape` on a Tracer with an UnshapedArray for its abstract value. Ugh!

We should probably make `onp.shape(unshaped)` raise a better error message, but the more important action item here is to raise an error message when the output of the custom jvp rule has inconsistent shapes for the primal and tangent values.

That's how JAX can raise a better error. But also we need to fix this code here so that both primal and tangent outputs are scalars (or both are 1D arrays with the same shape).

WDYT?
(I meant to add these kinds of error checks to #2026 a few weeks ago, but then had to stop working on it for a few weeks and forgot my context...)

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -708,6 +708,13 @@ def strip_weak_type(self):
     """Returns a copy of the aval with weak_type=False."""
     return UnshapedArray(self.dtype) if self.weak_type else self
 
+  @property
+  def shape(self):
+    msg = ("UnshapedArray has no shape. Please open an issue at "
+           "https://github.com/google/jax/issues because it's unexpected for "
+           "UnshapedArray instances to ever be produced.")
+    raise TypeError(msg)
+
 class ShapedArray(UnshapedArray):
   __slots__ = ['shape']
   array_abstraction_level = 1
diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -186,8 +186,23 @@ def _flatten_jvp(in_tree, *args):
   tangents_out, out_tree2 = tree_flatten(py_tangents_out)
   if out_tree != out_tree2:
     msg = ("Custom JVP rule must produce primal and tangent outputs with equal "
-           "container (pytree) structures, but got {} and {}.")
+           "container (pytree) structures, but got {} and {} respectively.")
     raise TypeError(msg.format(out_tree, out_tree2)) from None
+  primal_avals_out = [raise_to_shaped(core.get_aval(x)) for x in primals_out]
+  tangent_avals_out = [raise_to_shaped(core.get_aval(t)) for t in tangents_out]
+  if primal_avals_out != tangent_avals_out:
+    if len(primal_avals_out) == 1:
+      (av1,), (av2,) = primal_avals_out, tangent_avals_out
+      msg = ("Custom JVP rule must produce primal and tangent outputs with "
+             "equal shapes and dtypes, but got {} and {} respectively.")
+      raise TypeError(msg.format(av1.str_short(), av2.str_short()))
+    else:
+      msg = ("Custom JVP rule must produce primal and tangent outputs with "
+             "equal shapes and dtypes, but got:\n{}")
+      disagreements = (
+          "  primal {} for tangent {}".format(av1.str_short(), av2.str_short())
+          for av1, av2 in zip(primal_avals_out, tangent_avals_out) if av1 != av2)
+      raise TypeError(msg.format('\n'.join(disagreements)))
   yield primals_out + tangents_out, out_tree
 
 def _custom_deriv_call_bind(primitive, f, *args, **params):
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -1309,9 +1309,9 @@ def _check_tree(func_name, expected_name, actual_tree, expected_tree):
 def _check_tree_and_avals(what, tree1, avals1, tree2, avals2):
   """Raises TypeError if (tree1, avals1) does not match (tree2, avals2).
 
-  Corresponding `tree` and `avals` must match in the sense that the number of leaves in
-  `tree` must be equal to the length of `avals`.
-  `what` will be prepended to details of the mismatch in TypeError.
+  Corresponding `tree` and `avals` must match in the sense that the number of
+  leaves in `tree` must be equal to the length of `avals`. `what` will be
+  prepended to details of the mismatch in TypeError.
   """
   if tree1 != tree2:
     msg = ("{} must have same type structure, got {} and {}.")


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2228,7 +2228,7 @@ def test_vmap_axes(self):
   def test_pmap(self):
     raise unittest.SkipTest("TODO")  # TODO(mattjj): write test
 
-  def test_missing_jvp_rule_error(self):
+  def test_missing_jvp_rule_error_message(self):
     @api.custom_jvp
     def foo(x):
       return x ** 2
@@ -2246,7 +2246,7 @@ def foo(x):
         r"No JVP defined for custom_jvp function foo using defjvp.",
         lambda: api.grad(foo)(2.))
 
-  def test_jvp_rule_inconsistent_pytree_structures_error(self):
+  def test_jvp_rule_inconsistent_pytree_structures_error_message(self):
     @api.custom_jvp
     def f(x):
       return (x**2,)
@@ -2263,12 +2263,32 @@ def foo_jvp(primals, tangents):
         re.escape(
             "Custom JVP rule must produce primal and tangent outputs "
             "with equal container (pytree) structures, but got "
-            "{} and {}.".format(
+            "{} and {} respectively.".format(
                 tree_util.tree_structure((1,)),
                 tree_util.tree_structure([1, 2]))
         ),
         lambda: api.jvp(f, (2.,), (1.,)))
 
+  def test_primal_tangent_aval_disagreement_error_message(self):
+    @api.custom_jvp
+    def f(x):
+      return x ** 2
+
+    @f.defjvp
+    def foo_jvp(primals, tangents):
+      x, = primals
+      t, = tangents
+      return f(x), np.reshape(t, (1,))
+
+    f(2.)  # doesn't crash
+    self.assertRaisesRegex(
+        TypeError,
+        re.escape(
+            "Custom JVP rule must produce primal and tangent outputs "
+            "with equal shapes and dtypes, but got float32[] and float32[1] "
+            "respectively."),
+        lambda: api.jvp(f, (np.float32(2.),), (np.float32(1.),)))
+
 
 class CustomVJPTest(jtu.JaxTestCase):