mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 12:26:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Avoid creating float0 JVPTracers

Browse Source
This commit is contained in:
Dougal 2025-01-10 10:43:54 -05:00
parent c61b2f6b81
commit ba9b2ca5f6
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)