add custom_jvp / vjp, delete custom_transforms

fixes #116, #1097, #1249, #1275, #1366, #1723, #1670, #1875, #1938
jax-ml · Feb 12, 2020 · 0864ece · 0864ece
1 parent 5e77789
commit 0864ece
Show file tree

Hide file tree

Showing 15 changed files with 749 additions and 1,010 deletions.
diff --git a/docs/jax.rst b/docs/jax.rst
@@ -39,11 +39,6 @@ Automatic differentiation
 .. autofunction:: jvp
 .. autofunction:: linearize
 .. autofunction:: vjp
-.. autofunction:: custom_transforms
-.. autofunction:: defjvp
-.. autofunction:: defjvp_all
-.. autofunction:: defvjp
-.. autofunction:: defvjp_all
 .. autofunction:: custom_gradient
 
 

diff --git a/docs/notebooks/How_JAX_primitives_work.ipynb b/docs/notebooks/How_JAX_primitives_work.ipynb
@@ -1503,7 +1503,7 @@
             "  File \"/usr/local/lib/python3.6/dist-packages/jax/api.py\", line 611, in batched_fun\n",
             "    lambda: _flatten_axes(out_tree(), out_axes))\n",
             "  File \"/usr/local/lib/python3.6/dist-packages/jax/interpreters/batching.py\", line 41, in batch\n",
-            "    out_vals, out_dims = batch_fun(fun, in_vals, in_dims)\n",
+            "    out_vals, out_dims = batch2(fun, in_vals, in_dims)\n",
             "NotImplementedError: Batching rule for 'multiply_add' not implemented\n"
           ],
           "name": "stderr"