fix nabla for function with different in/out types (#4)

* fix nabla for function with different in/out types

* bump version

autofd/__init__.py

@@ -49,4 +49,4 @@
   "operators",
 ]
 
-__version__ = "0.0.4"  # noqa
+__version__ = "0.0.5"  # noqa

autofd/operators/operators.py

@@ -699,6 +699,9 @@ def nabla_f(*args):
   Returns:
     (nabla f): the gradient function.
   """
+  assert method in (
+    jax.jacfwd, jax.jacrev
+  ), "method need to be jax.jacfwd or jax.jacrev"
   return nabla_p.bind(f, argnums=argnums, method=method)
 
 
@@ -714,8 +717,11 @@ def nabla_spec(f, *, argnums, method):
   )
   ospec = f.ret.spec
   jspec = tree_map(
-    lambda osp:
-    tree_map(lambda isp: Spec(osp.shape + isp.shape, isp.dtype), ispec), ospec
+    lambda osp: tree_map(
+      lambda isp: Spec(
+        osp.shape + isp.shape, osp.dtype if method == jax.jacfwd else isp.dtype
+      ), ispec
+    ), ospec
   )
   return (Ret(jspec), *f.arg), None
 

setup.cfg

@@ -22,7 +22,7 @@ ignore = E731
 
 [metadata]
 name = autofd
-version = 0.0.4
+version = 0.0.5
 author = "Min Lin"
 author_email = "[email protected]"
 description = "Automatic Functional Derivative in JAX"