Add zero_from_primal

NeilGirdhar · Dec 23, 2024 · 62bcc19 · 62bcc19
1 parent f338a30
commit 62bcc19
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/tjax/__init__.py b/tjax/__init__.py
@@ -9,7 +9,8 @@
                                NumpyIntegralNumeric, NumpyRealArray, NumpyRealNumeric, PyTree,
                                RealArray, RealNumeric, Shape, ShapeLike, SliceLike)
 from ._src.cotangent_tools import (copy_cotangent, cotangent_combinator, print_cotangent,
-                                   replace_cotangent, reverse_scale_cotangent, scale_cotangent)
+                                   replace_cotangent, reverse_scale_cotangent, scale_cotangent,
+                                   zero_from_primal)
 from ._src.display.display_generic import display_generic
 from ._src.display.internal import internal_print_generic
 from ._src.display.print_generic import print_generic
@@ -116,4 +117,5 @@
     'scale_cotangent',
     'softplus',
     'tree_allclose',
+    'zero_from_primal',
 ]
diff --git a/tjax/_src/cotangent_tools.py b/tjax/_src/cotangent_tools.py
@@ -7,6 +7,7 @@
 
 import jax.numpy as jnp
 from jax import tree, vjp
+from jax.custom_derivatives import zero_from_primal as jax_zero_from_primal
 
 from .annotations import JaxRealArray, RealNumeric
 from .display.print_generic import print_generic
@@ -17,6 +18,10 @@
 Y = TypeVar('Y')
 
 
+def zero_from_primal(x: X, /, *, symbolic_zeros: bool = False) -> X:
+    return jax_zero_from_primal(x, symbolic_zeros=symbolic_zeros)
+
+
 # scale_cotangent ----------------------------------------------------------------------------------
 def scale_cotangent(x: X,
                     scalar_scale: RealNumeric | None = None,