Merge pull request #437 from google/issue433

fix symbolic zero handling in sub transpose
google · Feb 23, 2019 · aca4941 · aca4941
2 parents f08c3b7 + 1171626
commit aca4941
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/lax.py b/jax/lax.py
@@ -1363,7 +1363,7 @@ def _add_transpose(t, x, y):
 
 def _sub_transpose(t, x, y):
   assert x is None and y is None  # computation must be linear, not affine
-  return [t, neg(t)]
+  return [t, neg(t) if t is not ad_util.zero else ad_util.zero]
 
 sub_p = standard_binop([_num, _num], 'sub')
 ad.defjvp(sub_p,