jax-ml / jax-ml/jax

`Transpose rule for 'custom_vjp_call_jaxpr' not implemented` for `grad(custom_jvp(custom_vjp))`

Open
#19,087 4 comments 0 reactions 0 assignees View on GitHub
better_errors
Dominant language
Python
Stars
36.3k
Forks
3.8k
Avg merge
1d 6h
Merged PRs (30d)
399

Description

### Description

```python
import jax

@jax.custom_jvp
def f(x, y):
assert False

@f.defjvp
def f_jvp(primals, tangents):
x, y = primals
tx, ty = tangents
return x + y, g(tx, ty)

@jax.custom_vjp
def g(tx, ty):
return tx + ty

def g_fwd(tx, ty):
assert False

def g_bwd(res, ct):
assert False

g.defvjp(g_fwd, g_bwd)

jax.grad(f)(1., 1.)
```
produces
```python
File ".../jax/_src/interpreters/ad.py", line 287, in get_primitive_transpose
return primitive_transposes[p]
~~~~~~~~~~~~~~~~~~~~^^^
KeyError: custom_vjp_call_jaxpr

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

Traceback (most recent call last):
File "file.py", line 25, in
jax.grad(f)(1., 1.)
File "file.py", line 11, in f_jvp
return x + y, g(tx, ty)
jax._src.source_info_util.JaxStackTraceBeforeTransformation: NotImplementedError: Transpose rule (for reverse-mode differentiation) for 'custom_vjp_call_jaxpr' not implemented

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:

Traceback (most recent call last):
File "file.py", line 25, in
jax.grad(f)(1., 1.)
NotImplementedError: Transpose rule (for reverse-mode differentiation) for 'custom_vjp_call_jaxpr' not implemented
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
```

I recognise that this is maybe a little odd -- using a `custom_vjp` on a linear program -- so I'm mostly raising this as a need for a better error.

### What jax/jaxlib version are you using?

0.4.23

Contributor guide

Open the contributing guide

Research direction

Run the reproducer from the issue against the current JAX checkout, then read jax/_src/interpreters/ad.py around get_primitive_transpose and the custom differentiation handling. The issue is done when this nested custom_jvp/custom_vjp case produces a clearer, actionable error instead of the current missing transpose-rule traceback.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.