custom_jvp crash on my function

[mariogeiger]

I my code I have a function on which I want to define custom jvp because it is linear in his 3 arguments.
I got an error so I tried to make a minimal example that reproduce the problem I get.
Thanks for the help

Edit0: same error without constant tensor
Edit1: same error with multiplication instead of einsum
Edit2: same error with only 2 arguments

Here is my minimal code:

import jax
import jax.numpy as jnp

@jax.custom_jvp
def f(a, b):
    a = a.reshape(-1, 1)
    b = b.reshape(1, -1)
    return a * b

f.defjvps(
    lambda t, _, a, b: f(t, b),
    lambda t, _, a, b: f(a, t),
)

x = jnp.ones((3,))

jax.jacobian(f, 0)(x, x)

The error I get is:

---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation         Traceback (most recent call last)
~/anaconda3/lib/python3.8/runpy.py in _run_module_as_main(***failed resolving arguments***)
    193         sys.argv[0] = mod_spec.origin
--> 194     return _run_code(code, main_globals, None,
    195                      "__main__", mod_spec)

~/anaconda3/lib/python3.8/runpy.py in _run_code(***failed resolving arguments***)
     86                        __spec__ = mod_spec)
---> 87     exec(code, run_globals)
     88     return run_globals

~/anaconda3/lib/python3.8/site-packages/ipykernel_launcher.py in <module>
     15     from ipykernel import kernelapp as app
---> 16     app.launch_new_instance()

~/anaconda3/lib/python3.8/site-packages/traitlets/config/application.py in launch_instance(***failed resolving arguments***)
    844         app.initialize(argv)
--> 845         app.start()
    846 

~/anaconda3/lib/python3.8/site-packages/ipykernel/kernelapp.py in start(***failed resolving arguments***)
    611             try:
--> 612                 self.io_loop.start()
    613             except KeyboardInterrupt:

~/anaconda3/lib/python3.8/site-packages/tornado/platform/asyncio.py in start(***failed resolving arguments***)
    198             asyncio.set_event_loop(self.asyncio_loop)
--> 199             self.asyncio_loop.run_forever()
    200         finally:

~/anaconda3/lib/python3.8/asyncio/base_events.py in run_forever(***failed resolving arguments***)
    569             while True:
--> 570                 self._run_once()
    571                 if self._stopping:

~/anaconda3/lib/python3.8/asyncio/base_events.py in _run_once(***failed resolving arguments***)
   1858             else:
-> 1859                 handle._run()
   1860         handle = None  # Needed to break cycles when an exception occurs.

~/anaconda3/lib/python3.8/asyncio/events.py in _run(***failed resolving arguments***)
     80         try:
---> 81             self._context.run(self._callback, *self._args)
     82         except (SystemExit, KeyboardInterrupt):

~/anaconda3/lib/python3.8/site-packages/tornado/ioloop.py in <lambda>(***failed resolving arguments***)
    687             future.add_done_callback(
--> 688                 lambda f: self._run_callback(functools.partial(callback, future))
    689             )

~/anaconda3/lib/python3.8/site-packages/tornado/ioloop.py in _run_callback(***failed resolving arguments***)
    740         try:
--> 741             ret = callback()
    742             if ret is not None:

~/anaconda3/lib/python3.8/site-packages/tornado/gen.py in inner(***failed resolving arguments***)
    813                 f = None  # noqa: F841
--> 814                 self.ctx_run(self.run)
    815 

~/anaconda3/lib/python3.8/site-packages/tornado/gen.py in run(***failed resolving arguments***)
    774                     else:
--> 775                         yielded = self.gen.send(value)
    776 

~/anaconda3/lib/python3.8/site-packages/ipykernel/kernelbase.py in process_one(***failed resolving arguments***)
    364                 return None
--> 365         yield gen.maybe_future(dispatch(*args))
    366 

~/anaconda3/lib/python3.8/site-packages/tornado/gen.py in wrapper(***failed resolving arguments***)
    233                 try:
--> 234                     yielded = ctx_run(next, result)
    235                 except (StopIteration, Return) as e:

~/anaconda3/lib/python3.8/site-packages/ipykernel/kernelbase.py in dispatch_shell(***failed resolving arguments***)
    267             try:
--> 268                 yield gen.maybe_future(handler(stream, idents, msg))
    269             except Exception:

~/anaconda3/lib/python3.8/site-packages/tornado/gen.py in wrapper(***failed resolving arguments***)
    233                 try:
--> 234                     yielded = ctx_run(next, result)
    235                 except (StopIteration, Return) as e:

~/anaconda3/lib/python3.8/site-packages/ipykernel/kernelbase.py in execute_request(***failed resolving arguments***)
    542         reply_content = yield gen.maybe_future(
--> 543             self.do_execute(
    544                 code, silent, store_history,

~/anaconda3/lib/python3.8/site-packages/tornado/gen.py in wrapper(***failed resolving arguments***)
    233                 try:
--> 234                     yielded = ctx_run(next, result)
    235                 except (StopIteration, Return) as e:

~/anaconda3/lib/python3.8/site-packages/ipykernel/ipkernel.py in do_execute(***failed resolving arguments***)
    305                 # letting shell dispatch to loop runners
--> 306                 res = shell.run_cell(code, store_history=store_history, silent=silent)
    307         finally:

~/anaconda3/lib/python3.8/site-packages/ipykernel/zmqshell.py in run_cell(***failed resolving arguments***)
    535         self._last_traceback = None
--> 536         return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
    537 

~/anaconda3/lib/python3.8/site-packages/IPython/core/interactiveshell.py in run_cell(***failed resolving arguments***)
   2876         try:
-> 2877             result = self._run_cell(
   2878                 raw_cell, store_history, silent, shell_futures)

~/anaconda3/lib/python3.8/site-packages/IPython/core/interactiveshell.py in _run_cell(***failed resolving arguments***)
   2922         try:
-> 2923             return runner(coro)
   2924         except BaseException as e:

~/anaconda3/lib/python3.8/site-packages/IPython/core/async_helpers.py in _pseudo_sync_runner(***failed resolving arguments***)
     67     try:
---> 68         coro.send(None)
     69     except StopIteration as exc:

~/anaconda3/lib/python3.8/site-packages/IPython/core/interactiveshell.py in run_cell_async(***failed resolving arguments***)
   3145 
-> 3146                 has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
   3147                        interactivity=interactivity, compiler=compiler, result=result)

~/anaconda3/lib/python3.8/site-packages/IPython/core/interactiveshell.py in run_ast_nodes(***failed resolving arguments***)
   3337                         asy = compare(code)
-> 3338                     if (await self.run_code(code, result,  async_=asy)):
   3339                         return True

~/anaconda3/lib/python3.8/site-packages/IPython/core/interactiveshell.py in run_code(***failed resolving arguments***)
   3417                 else:
-> 3418                     exec(code_obj, self.user_global_ns, self.user_ns)
   3419             finally:

<ipython-input-3-af11faaa8a43> in <module>
     16 
---> 17 jax.jacobian(f, 0)(x, x)

<ipython-input-3-af11faaa8a43> in <lambda>(***failed resolving arguments***)
     10 f.defjvps(
---> 11     lambda t, _, a, b: f(t, b),
     12     lambda t, _, a, b: f(a, t),

<ipython-input-3-af11faaa8a43> in f(***failed resolving arguments***)
      7     b = b.reshape(1, -1)
----> 8     return a * b
      9 

~/anaconda3/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in deferring_binary_op(***failed resolving arguments***)
   5665       return NotImplemented
-> 5666     return binary_op(self, other)
   5667   return deferring_binary_op

~/anaconda3/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in fn(***failed resolving arguments***)
    426     x1, x2 = _promote_args(numpy_fn.__name__, x1, x2)
--> 427     return lax_fn(x1, x2) if x1.dtype != bool_ else bool_lax_fn(x1, x2)
    428   return _wraps(numpy_fn)(fn)

JaxStackTraceBeforeTransformation: AssertionError

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

UnfilteredStackTrace                      Traceback (most recent call last)
<ipython-input-3-af11faaa8a43> in <module>
     16 
---> 17 jax.jacobian(f, 0)(x, x)

~/anaconda3/lib/python3.8/site-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    142     try:
--> 143       return fun(*args, **kwargs)
    144     except Exception as e:

~/anaconda3/lib/python3.8/site-packages/jax/_src/api.py in batched_fun(*args, **kwargs)
   1298     axis_size = _mapped_axis_size(in_tree, args_flat, in_axes_flat, "vmap", kws=True)
-> 1299     out_flat = batching.batch(
   1300         flat_fun, axis_name, axis_size, in_axes_flat,

~/anaconda3/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    165     try:
--> 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:

~/anaconda3/lib/python3.8/site-packages/jax/_src/api.py in _vjp_pullback_wrapper(cotangent_dtypes, io_tree, fun, py_args)
   1885           f"with dtype {_dtype(arg)}.")
-> 1886   ans = fun(*args)
   1887   return tree_unflatten(out_tree, ans)

~/anaconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in unbound_vjp(pvals, jaxpr, consts, *cts)
    120     dummy_args = [UndefinedPrimal(v.aval) for v in jaxpr.invars]
--> 121     arg_cts = backward_pass(jaxpr, consts, dummy_args, cts)
    122     return map(instantiate_zeros, arg_cts)

~/anaconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in backward_pass(jaxpr, consts, primals_in, cotangents_in)
    217       else:
--> 218         cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals,
    219                                                          **eqn.params)

~/anaconda3/lib/python3.8/site-packages/jax/_src/custom_derivatives.py in _custom_jvp_call_jaxpr_transpose(***failed resolving arguments***)
    370   del jvp_jaxpr_thunk, num_consts
--> 371   return ad.backward_pass(fun_jaxpr.jaxpr, fun_jaxpr.consts, args, cts)
    372 ad.primitive_transposes[custom_jvp_call_jaxpr_p] = _custom_jvp_call_jaxpr_transpose

~/anaconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in backward_pass(jaxpr, consts, primals_in, cotangents_in)
    217       else:
--> 218         cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals,
    219                                                          **eqn.params)

~/anaconda3/lib/python3.8/site-packages/jax/_src/lax/lax.py in _mul_transpose(ct, x, y)
   2728 def _mul_transpose(ct, x, y):
-> 2729   assert ad.is_undefined_primal(x) ^ ad.is_undefined_primal(y)
   2730   if ad.is_undefined_primal(x):

UnfilteredStackTrace: AssertionError

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

AssertionError                            Traceback (most recent call last)
<ipython-input-3-af11faaa8a43> in <module>
     15 x = jnp.ones((3,))
     16 
---> 17 jax.jacobian(f, 0)(x, x)

~/anaconda3/lib/python3.8/site-packages/jax/_src/api.py in jacfun(*args, **kwargs)
   1048     y, pullback = _vjp(f_partial, *dyn_args)
   1049     tree_map(partial(_check_output_dtype_jacrev, holomorphic), y)
-> 1050     jac = vmap(pullback)(_std_basis(y))
   1051     jac = jac[0] if isinstance(argnums, int) else jac
   1052     example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args

~/anaconda3/lib/python3.8/site-packages/jax/_src/lax/lax.py in _mul_transpose(ct, x, y)
   2727 
   2728 def _mul_transpose(ct, x, y):
-> 2729   assert ad.is_undefined_primal(x) ^ ad.is_undefined_primal(y)
   2730   if ad.is_undefined_primal(x):
   2731     if type(ct) is ad_util.Zero:

AssertionError: 

[jakevdp]

Thanks for the report and for the clear repro - I think you're running into the issue mentioned in this comment: https://github.com/google/jax/blob/97a5719fcb40af7231b5f803f965063538282f8e/jax/interpreters/ad.py#L198-L200 This is a real bug and needs to be fixed.

One thing to note: the jax.jacobian function aliases to jax.jacrev, meaning that it uses reverse-mode autodiff, while your JVP implements forward-mode autodiff. JAX should be able to construct the backward rule automatically here, except for that bug in the backward pass above.

Until that is fixed, there are a couple possibilities:

• you may be able to limit your use-case to forward mode autodiff; i.e. call jax.jacfwd instead of jax.jacobian
• you could use jax.custom_vjp rather than jax.custom_jvp to implement reverse-mode autodiff manually

I opened a bug to track this in #7098

[mariogeiger]

Thanks for the kind answer. I should simply use jax.custom_vjp.

While thinking about this issue I had an idea. Would it be useful to allow the user to declare it's function linear in some argument, I'm thinking about something like that

@jax.declare_linear([0, 1])
def my_custom_scalar_product(x, y):
    ...

and maybe that way some operations like jax.grad but maybe also other operations could be implemented more efficiently?
We could imagine other form of declarations: declar_symmetric, declar_holomorphic, ...
In the same spirit the user could also be able to provide the inverse of a function or the adjoint of a linear function.
I'm curious to know if something like that already exists!