`AssertionError` with new `custom_jvp` for `jax.nn.relu`
On this [commit](https://github.com/google/jax/commit/ca23be63fbaed20192cda5f921afe177ac8dcf4d) (currently just a few commits behind head, none of which look relevant), I'm finding that
```python
import jax
import jax.numpy as np
import collections

class F(collections.namedtuple("F", ["a"])):
  def __call__(self, x):
    return jax.nn.relu(self.a) * x

@jax.jit
def g(f, x):
  return f(x)

jax.grad(F(2.0))(0.)  # works, returns DeviceArray(2., dtype=float32)
jax.grad(g, argnums=(1,))(F(2.0), 0.)  # AssertionError
```
hits an `AssertionError` on this line: https://github.com/google/jax/blob/1bb9aaa88c09a90a115c0304c05b0fc25932523b/jax/interpreters/partial_eval.py#L239

Commenting out the `@jax.jit` on `g` or replacing `jax.nn.relu(self.a)` with `np.maximum(self.a, 0)` both avoid the error.
diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
index bf8e9c5da..6e60bcd64 100644
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -252,6 +252,7 @@ def _flatten_jvp(in_tree, *args):
   yield primals_out + tangents_out, out_tree
 
 def _custom_jvp_call_bind(prim, fun, jvp, *args):
+  args = map(core.full_lower, args)
   top_trace = core.find_top_trace(args)
   level = (core.trace_state.trace_stack.next_level(True)
            if top_trace is None else top_trace.level)
@@ -490,6 +491,7 @@ def _flatten_bwd(in_tree, out_trees, *args):
   yield cts_in
 
 def _custom_vjp_call_bind(prim, fun, fwd, bwd, *args, out_trees):
+  args = map(core.full_lower, args)
   top_trace = core.find_top_trace(args)
   level = (core.trace_state.trace_stack.next_level(True)
            if top_trace is None else top_trace.level)
diff --git a/tests/api_test.py b/tests/api_test.py
index 365f4149e..ccdc5218c 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2781,6 +2781,19 @@ class CustomVJPTest(jtu.JaxTestCase):
     expected = jax.grad(f, 0)(2., 0.1) + jax.grad(f, 0)(2., 0.2)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def test_lowering_out_of_traces(self):
+    # https://github.com/google/jax/issues/2578
+
+    class F(collections.namedtuple("F", ["a"])):
+      def __call__(self, x):
+        return jax.nn.relu(self.a) * x
+
+    @jax.jit
+    def g(f, x):
+      return f(x)
+
+    jax.grad(g, argnums=(1,))(F(2.0), 0.)  # doesn't crash
+
 
 class DeprecatedCustomTransformsTest(jtu.JaxTestCase):
 
['jax/custom_derivatives.py', 'tests/api_test.py']