mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix custom_jvp check for tracers in arguments marked nondiff_argnums

Browse Source 
PiperOrigin-RevId: 520379098
This commit is contained in:
Roy Frostig 2023-03-29 10:41:34 -07:00 committed by jax authors
parent f282c251d4
commit c8a7d5990d
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/custom_derivatives.py
View File

@ -568,7 +568,7 @@ class custom_vjp(Generic[ReturnValue]):
def _check_for_tracers(x):
for leaf in tree_leaves(x):
if isinstance(x, core.Tracer):
if isinstance(leaf, core.Tracer):
msg = ("Found a JAX Tracer object passed as an argument to a custom_vjp "
"function in a position indicated by nondiff_argnums as "
"non-differentiable. Tracers cannot be passed as non-differentiable "