Neo wrapper

desc/derivatives.py

@@ -339,13 +339,13 @@ def f(x):
             tempargs = args[0 : self._argnum] + (x,) + args[self._argnum + 1 :]
             return self._fun(*tempargs)
 
-        x = np.atleast_1d(args[self._argnum])
+        x = jnp.atleast_1d(args[self._argnum])
         n = len(x)
         fx = f(x)
-        h = np.maximum(1.0, np.abs(x)) * self.rel_step
-        ee = np.diag(h)
+        h = jnp.maximum(1.0, np.abs(x)) * self.rel_step
+        ee = jnp.diag(h)
         dtype = fx.dtype
-        hess = np.outer(h, h)
+        hess = jnp.outer(h, h)
 
         for i in range(n):
             eei = ee[i, :]
@@ -383,13 +383,13 @@ def f(x):
             tempargs = args[0 : self._argnum] + (x,) + args[self._argnum + 1 :]
             return self._fun(*tempargs)
 
-        x0 = np.atleast_1d(args[self._argnum])
+        x0 = jnp.atleast_1d(args[self._argnum])
         f0 = f(x0)
         m = f0.size
         n = x0.size
-        J = np.zeros((m, n))
-        h = np.maximum(1.0, np.abs(x0)) * self.rel_step
-        h_vecs = np.diag(np.atleast_1d(h))
+        J = jnp.zeros((m, n))
+        h = jnp.maximum(1.0, np.abs(x0)) * self.rel_step
+        h_vecs = jnp.diag(np.atleast_1d(h))
         for i in range(n):
             x1 = x0 - h_vecs[i]
             x2 = x0 + h_vecs[i]
@@ -400,11 +400,46 @@ def f(x):
             dfdx = df / dx
             J = put(J.T, i, dfdx.flatten()).T
         if m == 1:
-            J = np.ravel(J)
+            J = jnp.ravel(J)
         return J
 
     @classmethod
     def compute_jvp(cls, fun, argnum, v, *args, **kwargs):
+        rel_step = kwargs.get("rel_step", 1e-3)
+        h = rel_step
+        if jnp.isscalar(argnum):
+            nargs = 1
+            argnum = (argnum,)
+        else:
+            nargs = len(
+                argnum,
+            )
+        v = (v,) if not isinstance(v, tuple) else v
+        v = v[:-1] + (jnp.array(v[-1]),)
+        args = args[:-1] + (jnp.array(args[-1]),)
+        varr, vtreedef, vidx = cls._tree2arr(v)
+        xarr, xtreedef, xidx = cls._tree2arr(args)
+
+        def fwrap(x):
+            x = cls._arr2tree(x, xtreedef, xidx)
+            return fun(*x)
+
+        fx = fwrap(xarr)
+        return (fwrap(xarr + h * varr) - fx) / h
+
+    @classmethod
+    def _tree2arr(cls, tree):
+        leaves, treedef = jax.tree_util.tree_flatten(tree)
+        idx = np.cumsum([foo.size for foo in leaves])
+        return jnp.concatenate(leaves), treedef, idx[:-1]
+
+    @classmethod
+    def _arr2tree(cls, arr, treedef, idx):
+        leaves = jnp.split(arr, idx)
+        return jax.tree_util.tree_unflatten(treedef, leaves)
+
+    @classmethod
+    def compute_jvp_old(cls, fun, argnum, v, *args, **kwargs):
         """Compute df/dx*v.
 
         Parameters
@@ -426,20 +461,19 @@ def compute_jvp(cls, fun, argnum, v, *args, **kwargs):
         """
         rel_step = kwargs.get("rel_step", 1e-3)
 
-        if np.isscalar(argnum):
+        if jnp.isscalar(argnum):
             nargs = 1
             argnum = (argnum,)
         else:
             nargs = len(argnum)
         v = (v,) if not isinstance(v, tuple) else v
-
-        f = np.array(
+        f = jnp.array(
             [
                 cls._compute_jvp_1arg(fun, argnum[i], v[i], *args, rel_step=rel_step)
                 for i in range(nargs)
             ]
         )
-        return np.sum(f, axis=0)
+        return jnp.sum(f, axis=0)
 
     @classmethod
     def compute_jvp2(cls, fun, argnum1, argnum2, v1, v2, *args):
@@ -538,11 +572,8 @@ def _compute_jvp(self, v, *args):
     def _compute_jvp_1arg(cls, fun, argnum, v, *args, **kwargs):
         """Compute a jvp wrt a single argument."""
         rel_step = kwargs.get("rel_step", 1e-3)
-        normv = np.linalg.norm(v)
-        if normv != 0:
-            vh = v / normv
-        else:
-            vh = v
+        normv = jnp.linalg.norm(v)
+        vh = jnp.where(normv != 0, v / normv, v)
         x = args[argnum]
 
         def f(x):