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

Merge pull request #25831 from jax-ml:avoid-float0-tracers

Browse Source 
PiperOrigin-RevId: 714058085
This commit is contained in:
jax authors 2025-01-10 08:16:43 -08:00
commit 4f106b8a27
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/interpreters/ad.py
View File

@ -511,7 +511,7 @@ class JVPTrace(Trace):
return map(partial(maybe_jvp_tracer, self), ps_out, ts_out)
def maybe_jvp_tracer(trace, primal, tangent):
if type(tangent) is Zero:
if type(tangent) is Zero or dtype(tangent) == float0:
return primal
else:
return JVPTracer(trace, primal, tangent)