`Transpose rule for 'custom_vjp_call_jaxpr' not implemented` for `grad(custom_jvp(custom_vjp))`
- 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
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