change ad_util const name (#9)

sail-sg · Feb 20, 2024 · 6284f32 · 6284f32
1 parent 1f1bd9c
commit 6284f32
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/autofd/__init__.py b/autofd/__init__.py
@@ -60,4 +60,4 @@ def scale_by_learning_rate(
   "operators",
 ]
 
-__version__ = "0.0.6"  # noqa
+__version__ = "0.0.7"  # noqa
diff --git a/autofd/operators/operators.py b/autofd/operators/operators.py
@@ -1196,8 +1196,8 @@ def add_values(*args):
     return fs[0]
 
 
-ad_util.jaxval_adders[types.FunctionType] = add
-ad_util.jaxval_adders[function] = add
+ad_util.raw_jaxval_adders[types.FunctionType] = add
+ad_util.raw_jaxval_adders[function] = add
 
 array_operators = {
   "neg": lambda x: numpy.negative(x),  # noqa

diff --git a/setup.cfg b/setup.cfg
@@ -22,7 +22,7 @@ ignore = E731
 
 [metadata]
 name = autofd
-version = 0.0.6
+version = 0.0.7
 author = "Min Lin"
 author_email = "[email protected]"
 description = "Automatic Functional Derivative in JAX"
@@ -38,7 +38,7 @@ classifiers =
 packages = find:
 python_requires = >=3.9
 install_requires =
-    jax>=0.4.16
+    jax>=0.4.24
     jaxtyping>=0.2.21
 
 [options.packages.find]