Update ad_support.py

brainpy · Nov 22, 2024 · 2d0d136 · 2d0d136
1 parent dba2438
commit 2d0d136
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/brainpy/_src/math/op_register/ad_support.py b/brainpy/_src/math/op_register/ad_support.py
@@ -41,9 +41,14 @@ def _standard_jvp(jvp_rules, primitive: Primitive, primals, tangents, **params):
       r = tuple(rule(t, *primals, **params))
       tangents_out.append(r)
       assert tree_util.tree_structure(r) == tree
-  return val_out, functools.reduce(_add_tangents,
+  try:
+    return val_out, functools.reduce(_add_tangents,
                                    tangents_out,
                                    tree_util.tree_map(lambda a: ad.Zero.from_primal_value(a), val_out))
+  except:
+    return val_out, functools.reduce(_add_tangents,
+                                     tangents_out,
+                                     tree_util.tree_map(lambda a: ad.Zero.from_value(a), val_out))
 
 
 def _add_tangents(xs, ys):