Repository: jax-ml/jax
Pull request: make custom_transforms handle pytrees, add api.defvjp
Before this PR, using `custom_transforms` (from #419) with functions that took tuples/lists/dicts (i.e. pytrees) as arguments would silently fail. Moreover, defining custom rules that took as input or returned as output any pytree values wouldn't work. Finally, `custom_transforms` functions didn't automatically support vmap.

This PR fixes all that!

This PR also improves the API by adding `defjvp` and `defvjp` (cf. #636) functions (and their variants) to api.py, which can be called directly on a `custom_transforms` function rather than having to dig out a reference to the underlying primitive. There's also a `custom_gradient` convenience wrapper, which works like `tf.custom_gradient`.

This API will likely be the main mechanism by which users will define custom gradients.

TODO:
- [x] add docstrings

In the future we should add a tutorial explanation of how to do this in an .md file, like [in Autograd's tutorial](https://github.com/HIPS/autograd/blob/master/docs/tutorial.md#extend-autograd-by-defining-your-own-primitives), at which point we can finally close #116.

It would also be good to add more tests, especially for error messages.

Problem statement:
Easy api for custom primitives and vjps
JAX supports custom primitives and vjps, just like Autograd did. Improvements:
1) add this to documentation
2) add a minimal example of this in the examples section
3) add a wrapper function if appropriate?

Hints from issue discussion:
Thanks. 

I am most concerned about adding wrappers of external libraries that have already been implemented with other forms of computing backend. We use a set of operators for particle mesh simulation and FFTs.

When we are talking about custom primitives, we also need to notify the auto-differ how to pick the tensor operators on the operands.  How is this currently done in JAX? 

I believe in autograd this was done with a vector space object that knows how to serialize any operand into a numpy array, after which numpy functions are used for inner products etc. This may not always be desirable -- e.g. if data has to be partitioned to several MPI ranks, then serialization to a single MPI rank is not even going to fit into the memory. We weren't able to use autograd due to this.

Another thing to worry about is whether these customized primitives support higher order differentiations. If the vjp function itself needs to be an external routine (not set of primitives) then higher order differentiation and auto-jvp are probably both broken? Is this a supported case?
Not sure if this is the right place: how can I define custom vmap primitives; in my case I am calling an external function that already supports batches and I want to vmap the code surrounding the call of this external primitive.
@jonasrauber that question is from a while ago, but the short answer is that `custom_transforms` (as in `from jax import custom_transforms`) is for doing this. To be improved and documented...
~~Just sketched out a custom VJPs API last night: https://gist.github.com/mattjj/2ba580930472e8e04c1759737268af92~~

~~The example there is trivial, and there's a bit more bookkeeping to be done to handle general code. But our initial thinking is that we can have a `defvjp` and a `defvjp_all`, to be used with `@custom_transforms`, where the former lets you specify a vjp function for each positional argument and the latter lets you specify a vjp for all arguments at once. (Maybe we can also provide a `defvjp_all_staged` if you want to compute some reduced residual information on the forward pass, rather than saving all the argument values.)~~

~~The funny bookkeeping in that gist is due to the fact that in JAX we usually don't specify VJPs (reverse-mode rules) directly, and instead only specify forward-mode rules; JAX generates reverse-mode autodiff through a composition of forward-mode, partial evaluation, and transposition transformations. But if you want to specify a VJP rule directly, that gist shows a trick to do it.~~

That was all a work-in-progress. We've got something better now!

Commit message: fix typo

Patch to review:
diff --git a/docs/jax.rst b/docs/jax.rst
index 857f54932c9e..00a9eade7440 100644
--- a/docs/jax.rst
+++ b/docs/jax.rst
@@ -19,6 +19,6 @@ Module contents
 ---------------
 
 .. automodule:: jax
-    :members: jit, disable_jit, grad, value_and_grad, vmap, pmap, jacfwd, jacrev, hessian, jvp, linearize, vjp, make_jaxpr, eval_shape
+    :members: jit, disable_jit, grad, value_and_grad, vmap, pmap, jacfwd, jacrev, hessian, jvp, linearize, vjp, make_jaxpr, eval_shape, custom_transforms, defjvp, defjvp2, defjvp_all, defvjp, defvjp2, defvjp_all, custom_gradient
     :undoc-members:
     :show-inheritance:
diff --git a/jax/api.py b/jax/api.py
index 7b847a41133b..39425b1cfa0c 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -26,6 +26,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import itertools
 import operator as op
 import os
@@ -38,6 +39,7 @@
 
 from . import core
 from . import linear_util as lu
+from . import ad_util
 from .core import pack, eval_jaxpr
 from .api_util import (pytree_fun_to_jaxtupletree_fun, pytree_to_jaxtupletree,
                        pytree_fun_to_flatjaxtuple_fun, apply_jaxtree_fun, wraps,
@@ -71,8 +73,8 @@ def jit(fun, static_argnums=()):
 
   Args:
     fun: Function to be jitted. Should be a pure function, as side-effects may
-      only be executed once. Its positional arguments and return value should be
-      arrays, scalars, or standard Python containers (tuple/list/dict) thereof.
+      only be executed once. Its arguments and return value should be arrays,
+      scalars, or (nested) standard Python containers (tuple/list/dict) thereof.
 
       Positional arguments indicated by `static_argnums` can be anything at all,
       provided they are hashable and have an equality operation defined. Static
@@ -952,25 +954,608 @@ def _valid_jaxtype(arg):
     return True
 
 
+class CustomTransformsFunction(object):
+  def __init__(self, fun, prim):
+    self.fun = fun
+    self.prim = prim
+    wraps(fun)(self)
+
+  def __repr__(self):
+    return '<jax.custom_transforms function {fun}>'.format(fun=self.__name__)
+
+  def __call__(self, *args, **kwargs):
+    jax_args, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
+    jax_kwargs, kwargs_tree = pytree_to_jaxtupletree(kwargs)
+    out_tree = lu.Store()
+    ans = self.prim.bind(jax_kwargs, *jax_args, kwargs_tree=kwargs_tree,
+                         in_trees=in_trees, out_tree=out_tree)
+    return build_tree(out_tree.val, ans)
+
 def custom_transforms(fun):
-  name = getattr(fun, '__name__', '<unnamed user primitive>')
+  """Wraps a function so that its transformation behavior can be controlled.
+
+  A primary use case of ``custom_transforms`` is defining custom VJP rules (aka
+  custom gradients) for a Python function, while still supporting other
+  transformations like ``jax.jit`` and ``jax.vmap``. Custom differentiation
+  rules can be supplied using the ``jax.defjvp`` and ``jax.defvjp`` functions.
+
+  JAX transforms Python functions by tracing which primitive operations (like
+  black-box numerical kernels) are applied to the function's input arguments to
+  produce its output, then transforming that recorded composition of primitives
+  using a table of transformation rules for each primitive. That is, for each
+  transformation (e.g. ``jax.jvp`` or ``jax.vmap``) there is a table of rules
+  that specifies how each primitive is transformed. When a new primitive is
+  added to the system, it also needs a rule for each transformation we might
+  want to apply to it.
+
+  But in some cases instead of introducing a new primitive and specifying every
+  transformation rule, we might start with a traceable/transformable Python
+  function and want to override just its VJP (i.e. give it a custom gradient).
+  More concretely, we might start with a Python function to which we can already
+  apply ``jax.jit``, but we want to change how ``jax.grad`` behaves on it. These
+  are the cases that ``custom_transforms`` handles.
+
+  The ``custom_transforms`` decorator wraps ``fun`` so that its transformation
+  behavior can be overridden, but not all transformation rules need to be
+  specified manually. For rules that aren't overridden, a default rule is used
+  for each transformation that relies on tracing ``fun``.
+
+  Args:
+    fun: a Python callable. Must be functionally pure. Its arguments and return
+      value should be arrays, scalars, or (nested) standard Python containers
+      (tuple/list/dict) thereof.
+
+  Returns:
+    A Python callable with the same input/output and transformation behavior as
+    ``fun``, but for which custom transformation rules can be supplied, e.g.
+    using ``jax.defvjp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  -5.4667816
+  >>> jax.defvjp(f, lambda g, x: g * x)
+  >>> print(jax.grad(f)(3.))
+  3.0
+  """
+  name = getattr(fun, '__name__', '<unnamed custom_transforms primitive>')
   fun_p = core.Primitive(name)
-  fun_p.def_impl(fun)
 
-  # generic transformation implementations that rely on traceability of `fun`
-  fun_p.def_abstract_eval(partial(pe.abstract_eval_fun, fun))
-  xla.translations[fun_p] = partial(xla.lower_fun, fun)
-  ad.primitive_jvps[fun_p] = partial(jvp, fun)
-  # TODO(mattjj): batching
+  def fun_impl(jax_kwargs, *jax_args, **params):
+    args = map(build_tree, params.pop('in_trees'), jax_args)
+    kwargs = build_tree(params.pop('kwargs_tree'), jax_kwargs)
+    pytree_out = fun(*args, **kwargs)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    params.pop('out_tree').store(out_tree)  # linear_util style side effect
+    assert not params
+    return out
+  fun_p.def_impl(fun_impl)
+
+  def fun_jvp(primals, tangents, **params):
+    return ad.jvp(lu.wrap_init(fun_impl, params)).call_wrapped(primals, tangents)
+  ad.primitive_jvps[fun_p] = fun_jvp
+
+  def fun_batch(batched_args, batch_dims, **params):
+    out = batching.batch(lu.wrap_init(fun_impl, params), batched_args, batch_dims, 0)
+    return out, 0
+  batching.primitive_batchers[fun_p] = fun_batch
+
+  staged_fun_p = core.Primitive('staged_' + name)
+  def fun_partial_eval(trace, *tracers, **params):
+    tracers = tuple(map(trace.instantiate_const, tracers))
+    avals = [t.aval for t in tracers]
+    pvals_in = [pe.PartialVal((a, core.unit)) for a in avals]
+    jaxpr, pval_out, consts = pe.trace_to_jaxpr(lu.wrap_init(fun_impl, params),
+                                                pvals_in, instantiate=True)
+    consts = trace.new_instantiated_const(core.pack(consts))
+    eqn = pe.JaxprEqn((consts,) + tracers, None, staged_fun_p, (), False, False,
+                      dict(params, jaxpr=jaxpr))
+    return pe.JaxprTracer(trace, pval_out, eqn)
+  pe.custom_partial_eval_rules[fun_p] = fun_partial_eval
+
+  def staged_fun_translation(c, xla_consts, *xla_args, **params):
+    consts_shapes = tuple(c.GetShape(xla_consts).tuple_shapes())
+    xla_consts = tuple(xla.xla_destructure(c, xla_consts))
+    arg_shapes = map(c.GetShape, xla_args)
+    built_c = xla.jaxpr_computation(params['jaxpr'], (), consts_shapes, *arg_shapes)
+    return c.Call(built_c, xla_consts + xla_args)
+  xla.translations[staged_fun_p] = staged_fun_translation
+
+  return CustomTransformsFunction(fun, fun_p)
+
+def _check_custom_transforms_type(name, fun):
+  if type(fun) is not CustomTransformsFunction:
+    msg = ("{} requires a custom_transforms function as its first argument, "
+          "but got type {}.")
+    raise TypeError(msg.format(name, type(fun)))
+
+def defjvp_all(fun, custom_jvp):
+  """Define a custom JVP rule for a ``custom_transforms`` function.
+
+  If ``fun`` represents a function with signature ``a -> b``, then
+  ``custom_jvp`` represents a function with signature ``a -> T a -> (b, Tb)``,
+  where we use ``T x`` to represent a tangent type for the type ``x``.
+
+  In more detail, ``custom_jvp`` must two arguments, both tuples of length equal
+  to the number of positional arguments to ``fun``. The first argument to
+  ``custom_jvp`` represents the input primal values, and the second represents
+  the input tangent values. ``custom_jvp`` must return a pair where the first
+  element represents the output primal value and the second element represents
+  the output tangent value.
+
+  Defining a custom JVP rule also affects the default VJP rule, which is derived
+  from the JVP rule automatically via transposition.
 
-  @wraps(fun)
-  def traceable(*args, **kwargs):
-    # TODO(mattjj): pytrees to jaxtupletrees
-    return fun_p.bind(*args, **kwargs)
-  traceable.primitive = fun_p
+  Args:
+    fun: a custom_transforms function.
+    custom_jvp: a Python callable specifying the JVP rule, taking two tuples as
+      arguments specifying the input primal values and tangent values,
+      respectively. The tuple elements can be arrays, scalars, or (nested)
+      standard Python containers (tuple/list/dict) thereof. Must be functionally
+      pure.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the JVP rule
+    specified by ``custom_jvp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  -10.933563
+  >>> jax.defjvp_all(f, lambda ps, ts: (np.sin(ps[0] ** 2), 8. * ts[0]))
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  16.0
+  """
+  _check_custom_transforms_type("defjvp_all", fun)
+  def custom_transforms_jvp(primals, tangents, **params):
+    jax_kwargs, jax_args = primals[0], primals[1:]
+    _, jax_args_dot = tangents[0], tangents[1:]
+    if jax_kwargs:
+      msg = ("defjvp_all requires the corresponding custom_transforms function "
+             "not to be called with keyword arguments.")
+      raise ValueError(msg)
+    in_trees = params['in_trees']
+    args = tuple(map(build_tree, in_trees, jax_args))
+    args_dot = tuple(map(build_tree, in_trees, jax_args_dot))
+    pytree_out, pytree_out_dot = custom_jvp(args, args_dot)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    out_dot, out_tree2 = pytree_to_jaxtupletree(pytree_out_dot)
+    if out_tree != out_tree2:
+      msg = ("custom jvp rule returned different tree structures for primals "
+             "and tangents, but they must be equal: {} vs {}.")
+      raise TypeError(msg.format(out_tree, out_tree2))
+    params['out_tree'].store(out_tree)  # linear_util style side effect
+    return out, out_dot
+  ad.primitive_jvps[fun.prim] = custom_transforms_jvp
+
+def defjvp(fun, *jvprules):
+  """Definine JVP rules for each argument separately.
+
+  This function is a convenience wrapper around ``jax.defjvp_all`` for
+  separately defining JVP rules for each of the function's arguments. This
+  convenience wrapper does not provide a mechanism for depending on anything
+  other than the function arguments, so JVP rules defined in this way can't
+  share work with the primal computation or with each other, though those things
+  are possible using ``jax.defjvp_all``. See ``jax.defjvp2`` for a variant that
+  allows for defining rules that depend on the output primal value. See also the
+  ``jax.defjvp_all`` docstring for more detail.
+
+  The signature of each component JVP rule is ``lambda g, *primals: ...`` where
+  ``g`` represents the tangent of the corresponding positional argument and
+  ``*primals`` represents all the primal positional arguments.
+
+  Defining a custom JVP rule also affects the default VJP rule, which is derived
+  from the JVP rule automatically via transposition.
+
+  Args:
+    fun: a custom_transforms function.
+    *jvprules: a sequence of functions or Nones specifying the JVP rule for each
+      corresponding positional argument. When an element is None, it indicates
+      that the Jacobian from the corresponding input to the output is zero.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the JVP rule
+    specified by ``*jvprules``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  -10.933563
+  >>> jax.defjvp(f, lambda g, x: 8. * g)
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  16.0
+
+  An example with a function on two arguments:
+
+  >>> @jax.custom_transforms
+  ... def f(x, y):
+  ...   return np.sin(x ** 2 + y)
+  ...
+  >>> print(f(3., 4.))
+  0.42016703
+  >>> out_primal, out_tangent = jax.jvp(f, (3., 4.), (1., 2.))
+  >>> print(out_primal)
+  0.42016703
+  >>> print(out_tangent)
+  7.2595744
+  >>> jax.defjvp(f, None, lambda g, x, y: 8. * g + y)
+  >>> out_primal, out_tangent = jax.jvp(f, (3., 4.), (1., 2.))
+  >>> print(out_primal)
+  0.42016703
+  >>> print(out_tangent)
+  20.0
+  """
+  _check_custom_transforms_type("defjvp", fun)
+  def custom_jvp(primals, tangents):
+    ans = fun(*primals)
+    tangents_out = [rule(t, *primals) for rule, t in zip(jvprules, tangents)
+                    if rule is not None and t is not ad_util.zero]
+    return ans, reduce(ad.add_tangents, tangents_out, ad_util.zero)
+  defjvp_all(fun, custom_jvp)
+
+def defjvp2(fun, *jvprules):
+  """Definine JVP rules for each argument separately.
+
+  This function is a convenience wrapper around ``jax.defjvp_all`` for separately
+  defining JVP rules for each of the function's arguments. This convenience
+  wrapper does not provide a mechanism for depending on anything other than the
+  function arguments and its primal output value.
+
+  The signature of each component JVP rule is ``lambda g, ans, *primals: ...``
+  where ``g`` represents the tangent of the corresponding positional argument,
+  ``ans`` represents the output primal, and ``*primals`` represents all the
+  primal positional arguments.
+
+  Defining a custom JVP rule also affects the default VJP rule, which is derived
+  from the JVP rule automatically via transposition.
 
-  return traceable
+  Args:
+    fun: a custom_transforms function.
+    *jvprules: a sequence of functions or Nones specifying the JVP rule for each
+      corresponding positional argument. When an element is None, it indicates
+      that the Jacobian from the corresponding input to the output is zero.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the JVP rule
+    specified by ``*jvprules``.
 
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  -10.933563
+  >>> jax.defjvp2(f, lambda g, ans, x: 8. * g + ans)
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  16.412119
+  """
+  _check_custom_transforms_type("defjvp2", fun)
+  def custom_jvp(primals, tangents):
+    ans = fun(*primals)
+    tangents_out = [rule(t, ans, *primals) for rule, t in zip(jvprules, tangents)
+                    if rule is not None and t is not ad_util.zero]
+    return ans, reduce(ad.add_tangents, tangents_out, ad_util.zero)
+  defjvp_all(fun, custom_jvp)
+
+def defvjp_all(fun, custom_vjp):
+  """Define a custom VJP rule for a ``custom_transforms`` function.
+
+  If ``fun`` represents a function with signature ``a -> b``, then
+  ``custom_vjp`` represents a function with signature ``a -> (b, CT b -> CT a)``
+  where we use ``CT x`` to represent a cotangent type for the type ``x``. That
+  is, ``custom_vjp`` should take the same arguments as ``fun`` and return a pair
+  where the first element represents the primal value of ``fun`` applied to the
+  arguments, and the second element is a VJP function that maps from output
+  cotangents to input cotangents, returning a tuple with length equal to the
+  number of positional arguments supplied to ``fun``.
+
+  The VJP function returned as the second element of the output of
+  ``custom_vjp`` can close over intermediate values computed when evaluating the
+  primal value of ``fun``. That is, use lexical closure to share work between
+  the forward pass and the backward pass of reverse-mode automatic
+  differentiation.
+
+  See also ``jax.custom_gradient``.
+
+  Args:
+    fun: a custom_transforms function.
+    custom_vjp: a Python callable specifying the VJP rule, taking the same
+      arguments as ``fun`` and returning a pair where the first elment is the
+      value of ``fun`` applied to the arguments and the second element is a
+      Python callable representing the VJP map from output cotangents to input
+      cotangents. The returned VJP function must accept a value with the same
+      shape as the value of ``fun`` applied to the arguments and must return a
+      tuple with length equal to the number of positional arguments to ``fun``.
+      Arguments can be arrays, scalars, or (nested) standard Python containers
+      (tuple/list/dict) thereof. Must be functionally pure.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the VJP rule
+    specified by ``custom_vjp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  -5.4667816
+  >>> jax.defvjp_all(f, lambda x: (np.sin(x ** 2), lambda g: (g * x,)))
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  3.0
+
+  An example with a function on two arguments, so that the VJP function must
+  return a tuple of length two:
+
+  >>> @jax.custom_transforms
+  ... def f(x, y):
+  ...   return x * y
+  ...
+  >>> jax.defvjp_all(f, lambda x, y: (x * y, lambda g: (y, x)))
+  >>> print(f(3., 4.))
+  12.0
+  >>> print(jax.grad(f, argnums=(0, 1))(3., 4.))
+  (4.0, 3.0)
+  """
+  _check_custom_transforms_type("defvjp_all", fun)
+  def custom_transforms_vjp(jax_kwargs, *jax_args, **params):
+    if jax_kwargs:
+      msg = ("defvjp_all requires the corresponding custom_transforms function "
+             "not to be called with keyword arguments.")
+      raise ValueError(msg)
+    args = map(build_tree, params['in_trees'], jax_args)
+    pytree_out, vjp_pytree = custom_vjp(*args)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    params['out_tree'].store(out_tree)  # linear_util style side effect
+    def vjp_pytree_(ct):
+      args_cts = tuple(vjp_pytree(ct))
+      if len(args_cts) != len(params['in_trees']):
+        msg = ("custom VJP function must return a tuple of length equal to the "
+               "number of positional arguments to the function being "
+               "differentiated: expected {}, got {}")
+        raise TypeError(msg.format(len(params['in_trees']), len(args_cts)))
+      return ({},) + args_cts
+    vjp, _ = pytree_fun_to_jaxtupletree_fun(lu.wrap_init(vjp_pytree_), (out_tree,))
+    return out, vjp.call_wrapped
+  ad.defvjp_all(fun.prim, custom_transforms_vjp)
+
+def defvjp(fun, *vjprules):
+  """Define VJP rules for each argument separately.
+
+  This function is a convenience wrapper around ``jax.defvjp_all`` for
+  separately defining VJP rules for each of the function's arguments. This
+  convenience wrapper does not provide a mechanism for depending on anything
+  other than the function arguments, so VJP rules defined in this way can't
+  share work with the primal computation or with each other, though those things
+  are possible using ``jax.defvjp_all`` or ``jax.custom_gradient``. See
+  ``jax.defvjp2`` for a variant that allows for defining rules that depend on
+  the output primal value. See also the ``jax.defvjp_all`` docstring.
+
+  The signature of each component VJP rule is ``lambda g, *primals: ...`` where
+  ``g`` represents the output cotangent and ``*primals`` represents all the
+  primal positional arguments.
+
+  Args:
+    fun: a custom_transforms function.
+    *vjprules: a sequence of functions or Nones specifying the VJP rule for each
+      corresponding positional argument. When an element is None, it indicates
+      that the Jacobian from the corresponding input to the output is zero.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the VJP rule
+    specified by ``*vjprules``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  >>> jax.defvjp(f, lambda g, x: 8. * g)
+  >>> print(jax.grad(f)(3.))
+  8.0
+
+  An example with a function on two arguments:
+
+  >>> @jax.custom_transforms
+  ... def f(x, y):
+  ...   return np.sin(x ** 2 + y)
+  ...
+  >>> print(f(3., 4.))
+  0.42016703
+  >>> print(jax.grad(f)(3., 4.))
+  5.4446807
+  >>> print(jax.grad(f, 1)(3., 4.))
+  0.9074468
+  >>> jax.defvjp(f, None, lambda g, x, y: g + x + y)
+  >>> print(jax.grad(f)(3., 4.))
+  0.0
+  >>> print(jax.grad(f, 1)(3., 4.))
+  8.0
+  """
+  _check_custom_transforms_type("defvjp", fun)
+  def custom_vjp(*primals):
+    ans = fun(*primals)
+    # TODO(mattjj): avoid instantiating zeros?
+    vjpfun = lambda ct: [vjp(ct, *primals) if vjp else ad_util.zeros_like_jaxval(x)
+                         for x, vjp in zip(primals, vjprules)]
+    return ans, vjpfun
+  defvjp_all(fun, custom_vjp)
+
+def defvjp2(fun, *vjprules):
+  """Define VJP rules for each argument separately.
+
+  This function is a convenience wrapper around ``jax.defvjp_all`` for
+  separately defining VJP rules for each of the function's arguments. This
+  convenience wrapper does not provide a mechanism for depending on anything
+  other than the function arguments and its primal output value.
+
+  The signature of each component VJP rule is ``lambda g, ans, *primals: ...``
+  where ``g`` represents the output cotangent, ``ans`` represents the output
+  primal, and ``*primals`` represents all the primal positional arguments.
+
+  Args:
+    fun: a custom_transforms function.
+    *vjprules: a sequence of functions or Nones specifying the VJP rule for each
+      corresponding positional argument. When an element is None, it indicates
+      that the Jacobian from the corresponding input to the output is zero.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the VJP rule
+    specified by ``*vjprules``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x, y):
+  ...   return np.sin(x ** 2 + y)
+  ...
+  >>> print(f(3., 4.))
+  0.42016703
+  >>> print(jax.grad(f)(3., 4.))
+  5.4446807
+  >>> print(jax.grad(f, 1)(3., 4.))
+  0.9074468
+  >>> jax.defvjp2(f, None, lambda g, ans, x, y: g + x + y + ans)
+  >>> print(jax.grad(f)(3., 4.))
+  0.0
+  >>> print(jax.grad(f, 1)(3., 4.))
+  8.420167
+  """
+  _check_custom_transforms_type("defvjp2", fun)
+  def custom_vjp(*primals):
+    ans = fun(*primals)
+    # TODO(mattjj): avoid instantiating zeros?
+    vjpfun = lambda ct: [vjp(ct, ans, *primals) if vjp else ad_util.zeros_like_jaxval(x)
+                         for x, vjp in zip(primals, vjprules)]
+    return ans, vjpfun
+  defvjp_all(fun, custom_vjp)
+
+def custom_gradient(fun):
+  """Convenience function for defining custom VJP rules (aka custom gradients).
+
+  While the canonical way to define custom VJP rules is via ``jax.defvjp_all``
+  and its convenience wrappers, the ``custom_gradient`` convenience wrapper
+  follows TensorFlow's tf.custom_gradient API.
+
+  See https://www.tensorflow.org/api_docs/python/tf/custom_gradient.
+
+  If the function to be differentiated has type signature ``a -> b``, then the
+  Python callable ``fun`` should have signature ``a -> (b, CT b -> CT a)`` where
+  we use ``CT x`` to denote a cotangent type for ``x``. See the example below.
+  That is, ``fun`` should return a pair where the first element represents the
+  value of the function to be differentiated and the second element is a
+  function that represents the custom VJP rule.
+
+  The custom VJP function returned as the second element of the output of ``fun``
+  can close over intermediate values computed when evaluating the function to be
+  differentiated. That is, use lexical closure to share work between the forward
+  pass and the backward pass of reverse-mode automatic differentiation.
+
+  Args:
+    fun: a Python callable specifying both the function to be differentiated and
+      its reverse-mode differentiation rule. It should return a pair consisting
+      of an output value and a Python callable that represents the custom
+      gradient function.
+
+  Returns:
+    A Python callable with signature ``a -> b``, i.e. that returns the output
+    value specified by the first element of ``fun``'s output pair. A side effect
+    is that under-the-hood ``jax.defvjp_all`` is called to set up the returned
+    Python callable with the custom VJP rule specified by the second element
+    of ``fun``'s output pair.
+
+  For example:
+
+  >>> @jax.custom_gradient
+  ... def f(x):
+  ...   return x ** 2, lambda g: (g * x,)
+  ...
+  >>> print(f(3.))
+  9.0
+  >>> print(jax.grad(f)(3.))
+  3.0
+
+  An example with a function on two arguments, so that the VJP function must
+  return a tuple of length two:
+
+  >>> @jax.custom_gradient
+  ... def f(x, y):
+  ...   return x * y, lambda g: (y, x)
+  ...
+  >>> print(f(3., 4.))
+  12.0
+  >>> print(jax.grad(f, argnums=(0, 1))(3., 4.))
+  (4.0, 3.0)
+  """
+  def primal_fun(*args, **kwargs):
+    ans, _ = fun(*args, **kwargs)
+    return ans
+  primal_fun = custom_transforms(primal_fun)
+  defvjp_all(primal_fun, fun)
+  return primal_fun
+
+
+def jarrett(fun):
+  new_fun = custom_transforms(fun)
+
+  def elementwise_jvp(primals, tangents):
+    pushfwd = partial(jvp, fun, primals)
+    y, jacs = vmap(pushfwd, out_axes=(None, 0))(_elementwise_std_basis(tangents))
+    flat_tangents, _ = tree_flatten(tangents)
+    out_tangent = sum([t * jac for t, jac in zip(flat_tangents, jacs)])
+    return y, out_tangent
+  defjvp_all(new_fun, elementwise_jvp)
+
+  return new_fun
 
 def _elementwise_std_basis(pytree):
   leaves, _ = tree_flatten(pytree)
@@ -987,19 +1572,6 @@ def _elementwise_std_basis(pytree):
        for j in range(arity)]) for i in range(arity)])
   return _unravel_array_into_pytree(pytree, 1, basis_array)
 
-def jarrett(fun):
-  new_fun = custom_transforms(fun)
-
-  def elementwise_jvp(primals, tangents):
-    pushfwd = partial(jvp, fun, primals)
-    y, jacs = vmap(pushfwd, out_axes=(None, 0))(_elementwise_std_basis(tangents))
-    flat_tangents, _ = tree_flatten(tangents)
-    out_tangent = sum([t * jac for t, jac in zip(flat_tangents, jacs)])
-    return y, out_tangent
-  ad.primitive_jvps[new_fun.primitive] = elementwise_jvp
-
-  return new_fun
-
 
 # This function mostly exists for making slides about JAX.
 def _make_graphviz(fun):
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
index 9d4e72fe68b0..428fb5c25094 100644
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -397,18 +397,19 @@ def add_tangents(x, y):
 
 
 def defvjp_all(prim, custom_vjp):
+  # see https://github.com/google/jax/pull/636
   name = prim.name
 
-  def fun_jvp(xs, ts):
+  def fun_jvp(xs, ts, **params):
     ts = map(instantiate_zeros, xs, ts)  # TODO(mattjj): avoid instantiation?
-    primal_out, tangent_out = fun_jvp_p.bind(pack(xs), pack(ts))
+    primal_out, tangent_out = fun_jvp_p.bind(pack(xs), pack(ts), **params)
     return primal_out, tangent_out
   primitive_jvps[prim] = fun_jvp
 
   fun_jvp_p = core.Primitive('{name}_jvp'.format(name=name))
-  def fun_jvp_partial_eval(trace, *tracers):
+  def fun_jvp_partial_eval(trace, *tracers, **params):
     primals_tracer, tangents_tracer = tracers
-    primal_out, vjp_py = custom_vjp(*primals_tracer)
+    primal_out, vjp_py = custom_vjp(*primals_tracer, **params)
 
     in_aval = raise_to_shaped(get_aval(primal_out))
     ct_pval = pe.PartialVal((in_aval, core.unit))
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
index 8295ebb39b6e..dc912221e176 100644
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -211,9 +211,10 @@ def partial_eval_wrapper(avals, *consts):
 
 
 def abstract_eval_fun(fun, *avals, **params):
-  pvs_in = [PartialVal((a, unit)) for a in avals]
-  _, pvout, _ = trace_to_jaxpr(lu.wrap_init(fun, params), pvs_in, instantiate=True)
-  aval_out, _ = pvout
+  pvals_in = [PartialVal((a, unit)) for a in avals]
+  _, pval_out, _ = trace_to_jaxpr(lu.wrap_init(fun, params), pvals_in,
+                                  instantiate=True)
+  aval_out, _ = pval_out
   assert isinstance(aval_out, AbstractValue)  # instantiate=True
   return aval_out
 
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
index dc6f5f023c33..3fa6e15745e9 100644
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -20,8 +20,7 @@
 import scipy.special as osp_special
 
 from .. import lax
-from ..api import custom_transforms
-from ..interpreters import ad, batching
+from ..api import custom_transforms, defjvp2
 from ..numpy import lax_numpy as np
 from ..numpy.lax_numpy import (_wraps, asarray, _reduction_dims, _constant_like,
                                _promote_args_like)
@@ -29,32 +28,32 @@
 
 @_wraps(osp_special.gammaln)
 def gammaln(x):
-    x, = _promote_args_like(osp_special.gammaln, x)
-    return lax.lgamma(x)
+  x, = _promote_args_like(osp_special.gammaln, x)
+  return lax.lgamma(x)
 
 
 @_wraps(osp_special.digamma)
 def digamma(x):
-    x, = _promote_args_like(osp_special.digamma, x)
-    return lax.digamma(x)
+  x, = _promote_args_like(osp_special.digamma, x)
+  return lax.digamma(x)
 
 
 @_wraps(osp_special.erf)
 def erf(x):
-    x, = _promote_args_like(osp_special.erf, x)
-    return lax.erf(x)
+  x, = _promote_args_like(osp_special.erf, x)
+  return lax.erf(x)
 
 
 @_wraps(osp_special.erfc)
 def erfc(x):
-    x, = _promote_args_like(osp_special.erfc, x)
-    return lax.erfc(x)
+  x, = _promote_args_like(osp_special.erfc, x)
+  return lax.erfc(x)
 
 
 @_wraps(osp_special.erfinv)
 def erfinv(x):
-    x, = _promote_args_like(osp_special.erfinv, x)
-    return lax.erf_inv(x)
+  x, = _promote_args_like(osp_special.erfinv, x)
+  return lax.erf_inv(x)
 
 
 @_wraps(osp_special.logit)
@@ -62,8 +61,7 @@ def erfinv(x):
 def logit(x):
   x = asarray(x)
   return lax.log(lax.div(x, lax.sub(lax._const(x, 1), x)))
-ad.defjvp2(logit.primitive, lambda g, ans, x: g / (x * (1 - x)))
-batching.defvectorized(logit.primitive)
+defjvp2(logit, lambda g, ans, x: g / (x * (1 - x)))
 
 
 @_wraps(osp_special.expit)
@@ -72,8 +70,7 @@ def expit(x):
   x = asarray(x)
   one = lax._const(x, 1)
   return lax.div(one, lax.add(one, lax.exp(lax.neg(x))))
-ad.defjvp2(expit.primitive, lambda g, ans, x: g * ans * (lax._const(ans, 1) - ans))
-batching.defvectorized(expit.primitive)
+defjvp2(expit, lambda g, ans, x: g * ans * (lax._const(ans, 1) - ans))
 
 
 @_wraps(osp_special.logsumexp)
diff --git a/tests/api_test.py b/tests/api_test.py
index 8c02fbd7ddf1..2b6507ea97ea 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -26,7 +26,7 @@
 from jax import jit, grad, device_get, device_put, jacfwd, jacrev, hessian
 from jax import api
 from jax.core import Primitive, pack, JaxTuple
-from jax.interpreters.ad import defjvp, defvjp, defvjp2, defvjp_all
+from jax.interpreters import ad
 from jax.interpreters.xla import DeviceArray, DeviceTuple
 from jax.abstract_arrays import concretization_err_msg
 from jax import test_util as jtu
@@ -242,7 +242,7 @@ def foo(x):
                      "XLA translation rule for 'foo' not implemented")
 
     foo_p.def_impl(lambda x: x)
-    defjvp(foo_p, lambda g, x: foo(g))
+    ad.defjvp(foo_p, lambda g, x: foo(g))
 
     jtu.check_raises(lambda: grad(foo)(1.0), NotImplementedError,
                      "Reverse-mode differentiation rule for 'foo' not implemented")
@@ -473,7 +473,7 @@ def test_defvjp_all(self):
     foo_p = Primitive('foo')
     def foo(x): return 2. * foo_p.bind(x)
 
-    defvjp_all(foo_p, lambda x: (x**2, lambda g: (4 * g * np.sin(x),)))
+    ad.defvjp_all(foo_p, lambda x: (x**2, lambda g: (4 * g * np.sin(x),)))
     val_ans, grad_ans = api.value_and_grad(foo)(3.)
     self.assertAllClose(val_ans, 2 * 3.**2, check_dtypes=False)
     self.assertAllClose(grad_ans, 4 * 2 * onp.sin(3.), check_dtypes=False)
@@ -482,7 +482,7 @@ def test_defvjp_all_const(self):
     foo_p = Primitive('foo')
     def foo(x): return foo_p.bind(x)
 
-    defvjp_all(foo_p, lambda x: (x**2, lambda g: (12.,)))
+    ad.defvjp_all(foo_p, lambda x: (x**2, lambda g: (12.,)))
     val_ans, grad_ans = api.value_and_grad(foo)(3.)
     self.assertAllClose(val_ans, 9., check_dtypes=False)
     self.assertAllClose(grad_ans, 12., check_dtypes=True)
@@ -491,7 +491,7 @@ def test_defvjp_all_higher_order_revmode(self):
     foo_p = Primitive('foo')
     def foo(x): return 2. * foo_p.bind(x)
 
-    defvjp_all(foo_p, lambda x: (x**2, lambda g: (g * x ** 2,)))
+    ad.defvjp_all(foo_p, lambda x: (x**2, lambda g: (g * x ** 2,)))
     ans = api.grad(api.grad(foo))(3.)
     self.assertAllClose(ans, 2 * 2 * 3., check_dtypes=False)
 
@@ -507,7 +507,7 @@ def vjpfun(x, y):
       vjp = lambda g: (g + x + y, g * x * 9.)
       return out, vjp
 
-    defvjp_all(foo_p, vjpfun)
+    ad.defvjp_all(foo_p, vjpfun)
     val_ans, grad_ans = api.value_and_grad(foo)(3., 4.)
     self.assertAllClose(val_ans, 3.**2 + 4.**3, check_dtypes=False)
     self.assertAllClose(grad_ans, 1. + 3. + 4., check_dtypes=False)
@@ -515,12 +515,22 @@ def vjpfun(x, y):
     ans = api.grad(foo, (0, 1))(3., 4.)
     self.assertAllClose(ans, (1. + 3. + 4., 1. * 3. * 9.), check_dtypes=False)
 
+  def test_defvjp_all(self):
+    @api.custom_transforms
+    def foo(x):
+      return np.sin(x)
+
+    api.defvjp_all(foo, lambda x: (np.sin(x), lambda g: (g * x,)))
+    val_ans, grad_ans = api.value_and_grad(foo)(3.)
+    self.assertAllClose(val_ans, onp.sin(3.), check_dtypes=False)
+    self.assertAllClose(grad_ans, 3., check_dtypes=False)
+
   def test_defvjp(self):
     @api.custom_transforms
     def foo(x, y):
       return np.sin(x * y)
 
-    defvjp(foo.primitive, None, lambda g, x, y: g * x * y)
+    api.defvjp(foo, None, lambda g, x, y: g * x * y)
     val_ans, grad_ans = api.value_and_grad(foo)(3., 4.)
     self.assertAllClose(val_ans, onp.sin(3. * 4.), check_dtypes=False)
     self.assertAllClose(grad_ans, 0., check_dtypes=False)
@@ -529,17 +539,87 @@ def foo(x, y):
     self.assertAllClose(ans_0, 0., check_dtypes=False)
     self.assertAllClose(ans_1, 3. * 4., check_dtypes=False)
 
+  def test_defvjp_higher_order(self):
+    @api.custom_transforms
+    def foo(x):
+      return np.sin(2. * x)
+
+    api.defvjp(foo, lambda g, x: g * np.cos(x))
+    ans = api.grad(api.grad(foo))(2.)
+    expected = api.grad(api.grad(np.sin))(2.)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   def test_defvjp2(self):
     @api.custom_transforms
     def foo(x, y):
       return np.sin(x * y)
 
-    defvjp2(foo.primitive, None, lambda g, ans, x, y: g * x * y + np.cos(ans))
+    api.defvjp2(foo, None, lambda g, ans, x, y: g * x * y + np.cos(ans))
     val_ans, grad_ans = api.value_and_grad(foo, 1)(3., 4.)
     self.assertAllClose(val_ans, onp.sin(3. * 4.), check_dtypes=False)
     self.assertAllClose(grad_ans, 3. * 4. + onp.cos(onp.sin(3. * 4)),
                         check_dtypes=False)
 
+  def test_custom_transforms_eval_with_pytrees(self):
+    @api.custom_transforms
+    def f(x):
+      a, b = x[0], x[1]
+      return {'hi': 2 * a, 'bye': 2 * b}
+
+    ans = f((1, 2))
+    self.assertEqual(ans, {'hi': 2 * 1, 'bye': 2 * 2})
+
+  def test_custom_transforms_jit_with_pytrees(self):
+    @api.custom_transforms
+    def f(x):
+      a, b = x[0], x[1]
+      return {'hi': 2 * a, 'bye': 2 * b}
+
+    ans = jit(f)((1, 2))
+    self.assertEqual(ans, {'hi': 2 * 1, 'bye': 2 * 2})
+
+  def test_custom_transforms_jit_with_pytrees_consts(self):
+    # The purpose of this test is to exercise the custom_transforms default
+    # translation rule in how it deals with constants that are too large to be
+    # treated as literals (at the time of writing).
+    z = onp.arange(10.)
+
+    @api.custom_transforms
+    def f(x):
+      a, b = x[0], x[1]
+      return {'hi': 2 * a, 'bye': z * b}
+
+    ans = jit(f)((1, 2))
+    self.assertAllClose(ans, {'hi': 2 * 1, 'bye': z * 2}, check_dtypes=False)
+
+  def test_custom_transforms_jvp_with_pytrees(self):
+    @api.custom_transforms
+    def f(x):
+      a, b = x[0], x[1]
+      return {'hi': 2 * a, 'bye': 2 * b}
+
+    ans, out_tangent = api.jvp(f, ((1, 2),), ((3, 4),))
+    self.assertEqual(ans, {'hi': 2 * 1, 'bye': 2 * 2})
+    self.assertEqual(out_tangent, {'hi': 2 * 3, 'bye': 2 * 4})
+
+  def test_custom_transforms_vmap_with_pytrees(self):
+    @api.custom_transforms
+    def f(x):
+      a, b = x[0], x[1]
+      return {'hi': 2 * a, 'bye': 2 * b}
+
+    ans = api.vmap(f)((onp.arange(3), onp.ones((3, 2))))
+    expected = {'hi': 2 * onp.arange(3), 'bye': 2 * onp.ones((3, 2))}
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def test_custom_gradient(self):
+    @api.custom_gradient
+    def f(x):
+      return x ** 2, lambda g: (g * x,)
+
+    self.assertEqual(f(3.), 9.)
+    self.assertEqual(api.grad(f)(3.), 3.)
+
   def test_devicetuple_iteration(self):
     tup = device_put(pack((1, 2)))
     self.assertIsInstance(tup, DeviceTuple)


Code review comments:
diff_hunk: @@ -952,25 +954,608 @@ def _valid_jaxtype(arg):
     return True
 
 
+class CustomTransformsFunction(object):
+  def __init__(self, fun, prim):
+    self.fun = fun
+    self.prim = prim
+    wraps(fun)(self)
+
+  def __repr__(self):
+    return '<jax.custom_transforms function {fun}>'.format(fun=self.__name__)
+
+  def __call__(self, *args, **kwargs):
+    jax_args, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
+    jax_kwargs, kwargs_tree = pytree_to_jaxtupletree(kwargs)
+    out_tree = lu.Store()
+    ans = self.prim.bind(jax_kwargs, *jax_args, kwargs_tree=kwargs_tree,
+                         in_trees=in_trees, out_tree=out_tree)
+    return build_tree(out_tree.val, ans)
+
 def custom_transforms(fun):
-  name = getattr(fun, '__name__', '<unnamed user primitive>')
+  """Wraps a function so that its transformation behavior can be controlled.
+
+  A primary use case of ``custom_transforms`` is defining custom VJP rules (aka
+  custom gradients) for a Python function, while still supporting other
+  transformations like ``jax.jit`` and ``jax.vmap``. Custom differentiation
+  rules can be supplied using the ``jax.defjvp`` and ``jax.defvjp`` functions.
+
+  JAX transforms Python functions by tracing which primitive operations (like
+  black-box numerical kernels) are applied to the function's input arguments to
+  produce its output, then transforming that recorded composition of primitives
+  using a table of transformation rules for each primitive. That is, for each
+  transformation (e.g. ``jax.jvp`` or ``jax.vmap``) there is a table of rules
+  that specifies how each primitive is transformed. When a new primitive is
+  added to the system, it also needs a rule for each transformation we might
+  want to apply to it.
+
+  But in some cases instead of introducing a new primitive and specifying every
+  transformation rule, we might start with a traceable/transformable Python
+  function and want to override just its VJP (i.e. give it a custom gradient).
+  More concretely, we might start with a Python function to which we can already
+  apply ``jax.jit``, but we want to change how ``jax.grad`` behaves on it. These
+  are the cases that ``custom_transforms`` handles.
path: jax/api.py
text: @user1:
I'm not sure the above two paragraphs are necessary. They're true and might be useful general info for understanding how jax works, but I don't think are needed to understand how to use this custom_transforms. This docstring is already quite long so I would save these paragraphs for a future in-depth guide :)

@author:
Done.

diff_hunk: @@ -952,25 +954,608 @@ def _valid_jaxtype(arg):
     return True
 
 
+class CustomTransformsFunction(object):
+  def __init__(self, fun, prim):
+    self.fun = fun
+    self.prim = prim
+    wraps(fun)(self)
+
+  def __repr__(self):
+    return '<jax.custom_transforms function {fun}>'.format(fun=self.__name__)
+
+  def __call__(self, *args, **kwargs):
+    jax_args, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
+    jax_kwargs, kwargs_tree = pytree_to_jaxtupletree(kwargs)
+    out_tree = lu.Store()
+    ans = self.prim.bind(jax_kwargs, *jax_args, kwargs_tree=kwargs_tree,
+                         in_trees=in_trees, out_tree=out_tree)
+    return build_tree(out_tree.val, ans)
+
 def custom_transforms(fun):
-  name = getattr(fun, '__name__', '<unnamed user primitive>')
+  """Wraps a function so that its transformation behavior can be controlled.
+
+  A primary use case of ``custom_transforms`` is defining custom VJP rules (aka
+  custom gradients) for a Python function, while still supporting other
+  transformations like ``jax.jit`` and ``jax.vmap``. Custom differentiation
+  rules can be supplied using the ``jax.defjvp`` and ``jax.defvjp`` functions.
+
+  JAX transforms Python functions by tracing which primitive operations (like
+  black-box numerical kernels) are applied to the function's input arguments to
+  produce its output, then transforming that recorded composition of primitives
+  using a table of transformation rules for each primitive. That is, for each
+  transformation (e.g. ``jax.jvp`` or ``jax.vmap``) there is a table of rules
+  that specifies how each primitive is transformed. When a new primitive is
+  added to the system, it also needs a rule for each transformation we might
+  want to apply to it.
+
+  But in some cases instead of introducing a new primitive and specifying every
+  transformation rule, we might start with a traceable/transformable Python
+  function and want to override just its VJP (i.e. give it a custom gradient).
+  More concretely, we might start with a Python function to which we can already
+  apply ``jax.jit``, but we want to change how ``jax.grad`` behaves on it. These
+  are the cases that ``custom_transforms`` handles.
+
+  The ``custom_transforms`` decorator wraps ``fun`` so that its transformation
+  behavior can be overridden, but not all transformation rules need to be
+  specified manually. For rules that aren't overridden, a default rule is used
+  for each transformation that relies on tracing ``fun``.
+
+  Args:
+    fun: a Python callable. Must be functionally pure. Its arguments and return
+      value should be arrays, scalars, or (nested) standard Python containers
+      (tuple/list/dict) thereof.
+
+  Returns:
+    A Python callable with the same input/output and transformation behavior as
+    ``fun``, but for which custom transformation rules can be supplied, e.g.
+    using ``jax.defvjp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  -5.4667816
+  >>> jax.defvjp(f, lambda g, x: g * x)
+  >>> print(jax.grad(f)(3.))
+  3.0
+  """
+  name = getattr(fun, '__name__', '<unnamed custom_transforms primitive>')
   fun_p = core.Primitive(name)
-  fun_p.def_impl(fun)
 
-  # generic transformation implementations that rely on traceability of `fun`
-  fun_p.def_abstract_eval(partial(pe.abstract_eval_fun, fun))
-  xla.translations[fun_p] = partial(xla.lower_fun, fun)
-  ad.primitive_jvps[fun_p] = partial(jvp, fun)
-  # TODO(mattjj): batching
+  def fun_impl(jax_kwargs, *jax_args, **params):
+    args = map(build_tree, params.pop('in_trees'), jax_args)
+    kwargs = build_tree(params.pop('kwargs_tree'), jax_kwargs)
+    pytree_out = fun(*args, **kwargs)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    params.pop('out_tree').store(out_tree)  # linear_util style side effect
+    assert not params
+    return out
+  fun_p.def_impl(fun_impl)
+
+  def fun_jvp(primals, tangents, **params):
+    return ad.jvp(lu.wrap_init(fun_impl, params)).call_wrapped(primals, tangents)
+  ad.primitive_jvps[fun_p] = fun_jvp
+
+  def fun_batch(batched_args, batch_dims, **params):
+    out = batching.batch(lu.wrap_init(fun_impl, params), batched_args, batch_dims, 0)
+    return out, 0
+  batching.primitive_batchers[fun_p] = fun_batch
+
+  staged_fun_p = core.Primitive('staged_' + name)
+  def fun_partial_eval(trace, *tracers, **params):
+    tracers = tuple(map(trace.instantiate_const, tracers))
+    avals = [t.aval for t in tracers]
+    pvals_in = [pe.PartialVal((a, core.unit)) for a in avals]
+    jaxpr, pval_out, consts = pe.trace_to_jaxpr(lu.wrap_init(fun_impl, params),
+                                                pvals_in, instantiate=True)
+    consts = trace.new_instantiated_const(core.pack(consts))
+    eqn = pe.JaxprEqn((consts,) + tracers, None, staged_fun_p, (), False, False,
+                      dict(params, jaxpr=jaxpr))
+    return pe.JaxprTracer(trace, pval_out, eqn)
+  pe.custom_partial_eval_rules[fun_p] = fun_partial_eval
+
+  def staged_fun_translation(c, xla_consts, *xla_args, **params):
+    consts_shapes = tuple(c.GetShape(xla_consts).tuple_shapes())
+    xla_consts = tuple(xla.xla_destructure(c, xla_consts))
+    arg_shapes = map(c.GetShape, xla_args)
+    built_c = xla.jaxpr_computation(params['jaxpr'], (), consts_shapes, *arg_shapes)
+    return c.Call(built_c, xla_consts + xla_args)
+  xla.translations[staged_fun_p] = staged_fun_translation
+
+  return CustomTransformsFunction(fun, fun_p)
+
+def _check_custom_transforms_type(name, fun):
+  if type(fun) is not CustomTransformsFunction:
+    msg = ("{} requires a custom_transforms function as its first argument, "
+          "but got type {}.")
+    raise TypeError(msg.format(name, type(fun)))
+
+def defjvp_all(fun, custom_jvp):
+  """Define a custom JVP rule for a ``custom_transforms`` function.
+
+  If ``fun`` represents a function with signature ``a -> b``, then
+  ``custom_jvp`` represents a function with signature ``a -> T a -> (b, Tb)``,
+  where we use ``T x`` to represent a tangent type for the type ``x``.
+
+  In more detail, ``custom_jvp`` must two arguments, both tuples of length equal
+  to the number of positional arguments to ``fun``. The first argument to
+  ``custom_jvp`` represents the input primal values, and the second represents
+  the input tangent values. ``custom_jvp`` must return a pair where the first
+  element represents the output primal value and the second element represents
+  the output tangent value.
+
+  Defining a custom JVP rule also affects the default VJP rule, which is derived
+  from the JVP rule automatically via transposition.
 
-  @wraps(fun)
-  def traceable(*args, **kwargs):
-    # TODO(mattjj): pytrees to jaxtupletrees
-    return fun_p.bind(*args, **kwargs)
-  traceable.primitive = fun_p
+  Args:
+    fun: a custom_transforms function.
+    custom_jvp: a Python callable specifying the JVP rule, taking two tuples as
+      arguments specifying the input primal values and tangent values,
+      respectively. The tuple elements can be arrays, scalars, or (nested)
+      standard Python containers (tuple/list/dict) thereof. Must be functionally
+      pure.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the JVP rule
+    specified by ``custom_jvp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  -10.933563
+  >>> jax.defjvp_all(f, lambda ps, ts: (np.sin(ps[0] ** 2), 8. * ts[0]))
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  16.0
+  """
+  _check_custom_transforms_type("defjvp_all", fun)
+  def custom_transforms_jvp(primals, tangents, **params):
+    jax_kwargs, jax_args = primals[0], primals[1:]
+    _, jax_args_dot = tangents[0], tangents[1:]
+    if jax_kwargs:
+      msg = ("defjvp_all requires the corresponding custom_transforms function "
+             "not to be called with keyword arguments.")
+      raise ValueError(msg)
+    in_trees = params['in_trees']
+    args = tuple(map(build_tree, in_trees, jax_args))
+    args_dot = tuple(map(build_tree, in_trees, jax_args_dot))
+    pytree_out, pytree_out_dot = custom_jvp(args, args_dot)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    out_dot, out_tree2 = pytree_to_jaxtupletree(pytree_out_dot)
+    if out_tree != out_tree2:
+      msg = ("custom jvp rule returned different tree structures for primals "
+             "and tangents, but they must be equal: {} vs {}.")
+      raise TypeError(msg.format(out_tree, out_tree2))
+    params['out_tree'].store(out_tree)  # linear_util style side effect
+    return out, out_dot
+  ad.primitive_jvps[fun.prim] = custom_transforms_jvp
+
+def defjvp(fun, *jvprules):
+  """Definine JVP rules for each argument separately.
+
+  This function is a convenience wrapper around ``jax.defjvp_all`` for
+  separately defining JVP rules for each of the function's arguments. This
+  convenience wrapper does not provide a mechanism for depending on anything
+  other than the function arguments, so JVP rules defined in this way can't
+  share work with the primal computation or with each other, though those things
+  are possible using ``jax.defjvp_all``. See ``jax.defjvp2`` for a variant that
+  allows for defining rules that depend on the output primal value. See also the
+  ``jax.defjvp_all`` docstring for more detail.
+
+  The signature of each component JVP rule is ``lambda g, *primals: ...`` where
+  ``g`` represents the tangent of the corresponding positional argument and
+  ``*primals`` represents all the primal positional arguments.
+
+  Defining a custom JVP rule also affects the default VJP rule, which is derived
+  from the JVP rule automatically via transposition.
+
+  Args:
+    fun: a custom_transforms function.
+    *jvprules: a sequence of functions or Nones specifying the JVP rule for each
+      corresponding positional argument. When an element is None, it indicates
+      that the Jacobian from the corresponding input to the output is zero.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the JVP rule
+    specified by ``*jvprules``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  -10.933563
+  >>> jax.defjvp(f, lambda g, x: 8. * g)
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  16.0
+
+  An example with a function on two arguments:
+
+  >>> @jax.custom_transforms
+  ... def f(x, y):
+  ...   return np.sin(x ** 2 + y)
+  ...
+  >>> print(f(3., 4.))
+  0.42016703
+  >>> out_primal, out_tangent = jax.jvp(f, (3., 4.), (1., 2.))
+  >>> print(out_primal)
+  0.42016703
+  >>> print(out_tangent)
+  7.2595744
+  >>> jax.defjvp(f, None, lambda g, x, y: 8. * g + y)
+  >>> out_primal, out_tangent = jax.jvp(f, (3., 4.), (1., 2.))
+  >>> print(out_primal)
+  0.42016703
+  >>> print(out_tangent)
+  20.0
+  """
+  _check_custom_transforms_type("defjvp", fun)
+  def custom_jvp(primals, tangents):
+    ans = fun(*primals)
+    tangents_out = [rule(t, *primals) for rule, t in zip(jvprules, tangents)
+                    if rule is not None and t is not ad_util.zero]
+    return ans, reduce(ad.add_tangents, tangents_out, ad_util.zero)
+  defjvp_all(fun, custom_jvp)
+
+def defjvp2(fun, *jvprules):
path: jax/api.py
text: @user1:
I'm not sure what 2 refers to, maybe the two non-variadic arguments to the jvprules functions? Could we get rid of defjvp and only have this one?

@author:
We could just have one, but the most common pattern in our internal code (where we have `ad.defjvp` and `ad.defjvp2` for Primitives) is not to depend on the output. That is, if you look at our JVP rules in lax.py, we almost always use ``defejvp`` but then sometimes use ``defjvp2``. That and having these functions follow the convention of their counterparts in ad.py seem like good reasons.

The 2 doesn't refer to anything, it's just a way to name a variant of a function while keeping the name short. To me it's like a "prime", but we can't add `'` to a name in Python.

I know you might not like that convention, but we already use it elsewhere (to no ill effect so far!), and in any case it doesn't seem like something worth getting stuck on. WDYT?

@user2:
Maybe `defjvp_separate` or  `devjvp_sep` if you really want to save characters?

For what it's worth I also found the name confusing, and wondered what the 2 was about :)

@user3:
I vote for the current name: `defjvp2` by the same reason as Matt.

@user1:
For jvp rules that don't depend on the output, can't you just put a `_` for that arg and ignore it? If the performance etc. is the same, seems pretty lightweight to me. People (including us) can make their own wrappers if it really bothers them.

@author:
> For what it's worth I also found the name confusing, and wondered what the 2 was about :)

Wondering what a function is and how it works doesn't seem like a bad thing! I would wonder about what all of these functions mean before looking at example uses or their code (or docstring). ("separate" isn't the right suffix here because the difference in question is between taking the output primal as an argument or not.)

> People (including us) can make their own wrappers if it really bothers them.

Yes, that's the idea, but since `defjvp` is the most common pattern, and `defjvp2` is the more general second-most-common pattern, we provide those wrappers for users. After all, both `defjvp` and `defjvp2` are already convenience wrappers around `defjvp_all`.

I think reasonable people can (and do) disagree about personal preferences here. I'm going to keep `defjvp2`.

diff_hunk: @@ -952,25 +954,608 @@ def _valid_jaxtype(arg):
     return True
 
 
+class CustomTransformsFunction(object):
+  def __init__(self, fun, prim):
+    self.fun = fun
+    self.prim = prim
+    wraps(fun)(self)
+
+  def __repr__(self):
+    return '<jax.custom_transforms function {fun}>'.format(fun=self.__name__)
+
+  def __call__(self, *args, **kwargs):
+    jax_args, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
+    jax_kwargs, kwargs_tree = pytree_to_jaxtupletree(kwargs)
+    out_tree = lu.Store()
+    ans = self.prim.bind(jax_kwargs, *jax_args, kwargs_tree=kwargs_tree,
+                         in_trees=in_trees, out_tree=out_tree)
+    return build_tree(out_tree.val, ans)
+
 def custom_transforms(fun):
-  name = getattr(fun, '__name__', '<unnamed user primitive>')
+  """Wraps a function so that its transformation behavior can be controlled.
+
+  A primary use case of ``custom_transforms`` is defining custom VJP rules (aka
+  custom gradients) for a Python function, while still supporting other
+  transformations like ``jax.jit`` and ``jax.vmap``. Custom differentiation
+  rules can be supplied using the ``jax.defjvp`` and ``jax.defvjp`` functions.
+
+  JAX transforms Python functions by tracing which primitive operations (like
+  black-box numerical kernels) are applied to the function's input arguments to
+  produce its output, then transforming that recorded composition of primitives
+  using a table of transformation rules for each primitive. That is, for each
+  transformation (e.g. ``jax.jvp`` or ``jax.vmap``) there is a table of rules
+  that specifies how each primitive is transformed. When a new primitive is
+  added to the system, it also needs a rule for each transformation we might
+  want to apply to it.
+
+  But in some cases instead of introducing a new primitive and specifying every
+  transformation rule, we might start with a traceable/transformable Python
+  function and want to override just its VJP (i.e. give it a custom gradient).
+  More concretely, we might start with a Python function to which we can already
+  apply ``jax.jit``, but we want to change how ``jax.grad`` behaves on it. These
+  are the cases that ``custom_transforms`` handles.
+
+  The ``custom_transforms`` decorator wraps ``fun`` so that its transformation
+  behavior can be overridden, but not all transformation rules need to be
+  specified manually. For rules that aren't overridden, a default rule is used
+  for each transformation that relies on tracing ``fun``.
+
+  Args:
+    fun: a Python callable. Must be functionally pure. Its arguments and return
+      value should be arrays, scalars, or (nested) standard Python containers
+      (tuple/list/dict) thereof.
+
+  Returns:
+    A Python callable with the same input/output and transformation behavior as
+    ``fun``, but for which custom transformation rules can be supplied, e.g.
+    using ``jax.defvjp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  -5.4667816
+  >>> jax.defvjp(f, lambda g, x: g * x)
+  >>> print(jax.grad(f)(3.))
+  3.0
+  """
+  name = getattr(fun, '__name__', '<unnamed custom_transforms primitive>')
   fun_p = core.Primitive(name)
-  fun_p.def_impl(fun)
 
-  # generic transformation implementations that rely on traceability of `fun`
-  fun_p.def_abstract_eval(partial(pe.abstract_eval_fun, fun))
-  xla.translations[fun_p] = partial(xla.lower_fun, fun)
-  ad.primitive_jvps[fun_p] = partial(jvp, fun)
-  # TODO(mattjj): batching
+  def fun_impl(jax_kwargs, *jax_args, **params):
+    args = map(build_tree, params.pop('in_trees'), jax_args)
+    kwargs = build_tree(params.pop('kwargs_tree'), jax_kwargs)
+    pytree_out = fun(*args, **kwargs)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    params.pop('out_tree').store(out_tree)  # linear_util style side effect
+    assert not params
+    return out
+  fun_p.def_impl(fun_impl)
+
+  def fun_jvp(primals, tangents, **params):
+    return ad.jvp(lu.wrap_init(fun_impl, params)).call_wrapped(primals, tangents)
+  ad.primitive_jvps[fun_p] = fun_jvp
+
+  def fun_batch(batched_args, batch_dims, **params):
+    out = batching.batch(lu.wrap_init(fun_impl, params), batched_args, batch_dims, 0)
+    return out, 0
+  batching.primitive_batchers[fun_p] = fun_batch
+
+  staged_fun_p = core.Primitive('staged_' + name)
+  def fun_partial_eval(trace, *tracers, **params):
+    tracers = tuple(map(trace.instantiate_const, tracers))
+    avals = [t.aval for t in tracers]
+    pvals_in = [pe.PartialVal((a, core.unit)) for a in avals]
+    jaxpr, pval_out, consts = pe.trace_to_jaxpr(lu.wrap_init(fun_impl, params),
+                                                pvals_in, instantiate=True)
+    consts = trace.new_instantiated_const(core.pack(consts))
+    eqn = pe.JaxprEqn((consts,) + tracers, None, staged_fun_p, (), False, False,
+                      dict(params, jaxpr=jaxpr))
+    return pe.JaxprTracer(trace, pval_out, eqn)
+  pe.custom_partial_eval_rules[fun_p] = fun_partial_eval
+
+  def staged_fun_translation(c, xla_consts, *xla_args, **params):
+    consts_shapes = tuple(c.GetShape(xla_consts).tuple_shapes())
+    xla_consts = tuple(xla.xla_destructure(c, xla_consts))
+    arg_shapes = map(c.GetShape, xla_args)
+    built_c = xla.jaxpr_computation(params['jaxpr'], (), consts_shapes, *arg_shapes)
+    return c.Call(built_c, xla_consts + xla_args)
+  xla.translations[staged_fun_p] = staged_fun_translation
+
+  return CustomTransformsFunction(fun, fun_p)
+
+def _check_custom_transforms_type(name, fun):
+  if type(fun) is not CustomTransformsFunction:
+    msg = ("{} requires a custom_transforms function as its first argument, "
+          "but got type {}.")
+    raise TypeError(msg.format(name, type(fun)))
+
+def defjvp_all(fun, custom_jvp):
+  """Define a custom JVP rule for a ``custom_transforms`` function.
+
+  If ``fun`` represents a function with signature ``a -> b``, then
+  ``custom_jvp`` represents a function with signature ``a -> T a -> (b, Tb)``,
+  where we use ``T x`` to represent a tangent type for the type ``x``.
+
+  In more detail, ``custom_jvp`` must two arguments, both tuples of length equal
+  to the number of positional arguments to ``fun``. The first argument to
+  ``custom_jvp`` represents the input primal values, and the second represents
+  the input tangent values. ``custom_jvp`` must return a pair where the first
+  element represents the output primal value and the second element represents
+  the output tangent value.
+
+  Defining a custom JVP rule also affects the default VJP rule, which is derived
+  from the JVP rule automatically via transposition.
 
-  @wraps(fun)
-  def traceable(*args, **kwargs):
-    # TODO(mattjj): pytrees to jaxtupletrees
-    return fun_p.bind(*args, **kwargs)
-  traceable.primitive = fun_p
+  Args:
+    fun: a custom_transforms function.
+    custom_jvp: a Python callable specifying the JVP rule, taking two tuples as
+      arguments specifying the input primal values and tangent values,
+      respectively. The tuple elements can be arrays, scalars, or (nested)
+      standard Python containers (tuple/list/dict) thereof. Must be functionally
+      pure.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the JVP rule
+    specified by ``custom_jvp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  -10.933563
+  >>> jax.defjvp_all(f, lambda ps, ts: (np.sin(ps[0] ** 2), 8. * ts[0]))
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  16.0
+  """
+  _check_custom_transforms_type("defjvp_all", fun)
+  def custom_transforms_jvp(primals, tangents, **params):
+    jax_kwargs, jax_args = primals[0], primals[1:]
+    _, jax_args_dot = tangents[0], tangents[1:]
+    if jax_kwargs:
+      msg = ("defjvp_all requires the corresponding custom_transforms function "
+             "not to be called with keyword arguments.")
+      raise ValueError(msg)
+    in_trees = params['in_trees']
+    args = tuple(map(build_tree, in_trees, jax_args))
+    args_dot = tuple(map(build_tree, in_trees, jax_args_dot))
+    pytree_out, pytree_out_dot = custom_jvp(args, args_dot)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    out_dot, out_tree2 = pytree_to_jaxtupletree(pytree_out_dot)
+    if out_tree != out_tree2:
+      msg = ("custom jvp rule returned different tree structures for primals "
+             "and tangents, but they must be equal: {} vs {}.")
+      raise TypeError(msg.format(out_tree, out_tree2))
+    params['out_tree'].store(out_tree)  # linear_util style side effect
+    return out, out_dot
+  ad.primitive_jvps[fun.prim] = custom_transforms_jvp
+
+def defjvp(fun, *jvprules):
+  """Definine JVP rules for each argument separately.
+
+  This function is a convenience wrapper around ``jax.defjvp_all`` for
+  separately defining JVP rules for each of the function's arguments. This
+  convenience wrapper does not provide a mechanism for depending on anything
+  other than the function arguments, so JVP rules defined in this way can't
+  share work with the primal computation or with each other, though those things
+  are possible using ``jax.defjvp_all``. See ``jax.defjvp2`` for a variant that
+  allows for defining rules that depend on the output primal value. See also the
+  ``jax.defjvp_all`` docstring for more detail.
+
+  The signature of each component JVP rule is ``lambda g, *primals: ...`` where
+  ``g`` represents the tangent of the corresponding positional argument and
+  ``*primals`` represents all the primal positional arguments.
+
+  Defining a custom JVP rule also affects the default VJP rule, which is derived
+  from the JVP rule automatically via transposition.
+
+  Args:
+    fun: a custom_transforms function.
+    *jvprules: a sequence of functions or Nones specifying the JVP rule for each
+      corresponding positional argument. When an element is None, it indicates
+      that the Jacobian from the corresponding input to the output is zero.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the JVP rule
+    specified by ``*jvprules``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  -10.933563
+  >>> jax.defjvp(f, lambda g, x: 8. * g)
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  16.0
+
+  An example with a function on two arguments:
+
+  >>> @jax.custom_transforms
+  ... def f(x, y):
+  ...   return np.sin(x ** 2 + y)
+  ...
+  >>> print(f(3., 4.))
+  0.42016703
+  >>> out_primal, out_tangent = jax.jvp(f, (3., 4.), (1., 2.))
+  >>> print(out_primal)
+  0.42016703
+  >>> print(out_tangent)
+  7.2595744
+  >>> jax.defjvp(f, None, lambda g, x, y: 8. * g + y)
+  >>> out_primal, out_tangent = jax.jvp(f, (3., 4.), (1., 2.))
+  >>> print(out_primal)
+  0.42016703
+  >>> print(out_tangent)
+  20.0
+  """
+  _check_custom_transforms_type("defjvp", fun)
+  def custom_jvp(primals, tangents):
+    ans = fun(*primals)
+    tangents_out = [rule(t, *primals) for rule, t in zip(jvprules, tangents)
+                    if rule is not None and t is not ad_util.zero]
+    return ans, reduce(ad.add_tangents, tangents_out, ad_util.zero)
+  defjvp_all(fun, custom_jvp)
+
+def defjvp2(fun, *jvprules):
+  """Definine JVP rules for each argument separately.
+
+  This function is a convenience wrapper around ``jax.defjvp_all`` for separately
+  defining JVP rules for each of the function's arguments. This convenience
+  wrapper does not provide a mechanism for depending on anything other than the
+  function arguments and its primal output value.
+
+  The signature of each component JVP rule is ``lambda g, ans, *primals: ...``
+  where ``g`` represents the tangent of the corresponding positional argument,
+  ``ans`` represents the output primal, and ``*primals`` represents all the
+  primal positional arguments.
+
+  Defining a custom JVP rule also affects the default VJP rule, which is derived
+  from the JVP rule automatically via transposition.
 
-  return traceable
+  Args:
+    fun: a custom_transforms function.
+    *jvprules: a sequence of functions or Nones specifying the JVP rule for each
+      corresponding positional argument. When an element is None, it indicates
+      that the Jacobian from the corresponding input to the output is zero.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the JVP rule
+    specified by ``*jvprules``.
 
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  -10.933563
+  >>> jax.defjvp2(f, lambda g, ans, x: 8. * g + ans)
+  >>> out_primal, out_tangent = jax.jvp(f, (3.,), (2.,))
+  >>> print(out_primal)
+  0.4121185
+  >>> print(out_tangent)
+  16.412119
+  """
+  _check_custom_transforms_type("defjvp2", fun)
+  def custom_jvp(primals, tangents):
+    ans = fun(*primals)
+    tangents_out = [rule(t, ans, *primals) for rule, t in zip(jvprules, tangents)
+                    if rule is not None and t is not ad_util.zero]
+    return ans, reduce(ad.add_tangents, tangents_out, ad_util.zero)
+  defjvp_all(fun, custom_jvp)
+
+def defvjp_all(fun, custom_vjp):
+  """Define a custom VJP rule for a ``custom_transforms`` function.
+
+  If ``fun`` represents a function with signature ``a -> b``, then
+  ``custom_vjp`` represents a function with signature ``a -> (b, CT b -> CT a)``
+  where we use ``CT x`` to represent a cotangent type for the type ``x``. That
+  is, ``custom_vjp`` should take the same arguments as ``fun`` and return a pair
+  where the first element represents the primal value of ``fun`` applied to the
+  arguments, and the second element is a VJP function that maps from output
+  cotangents to input cotangents, returning a tuple with length equal to the
+  number of positional arguments supplied to ``fun``.
+
+  The VJP function returned as the second element of the output of
+  ``custom_vjp`` can close over intermediate values computed when evaluating the
+  primal value of ``fun``. That is, use lexical closure to share work between
+  the forward pass and the backward pass of reverse-mode automatic
+  differentiation.
+
+  See also ``jax.custom_gradient``.
+
+  Args:
+    fun: a custom_transforms function.
+    custom_vjp: a Python callable specifying the VJP rule, taking the same
+      arguments as ``fun`` and returning a pair where the first elment is the
+      value of ``fun`` applied to the arguments and the second element is a
+      Python callable representing the VJP map from output cotangents to input
+      cotangents. The returned VJP function must accept a value with the same
+      shape as the value of ``fun`` applied to the arguments and must return a
+      tuple with length equal to the number of positional arguments to ``fun``.
+      Arguments can be arrays, scalars, or (nested) standard Python containers
+      (tuple/list/dict) thereof. Must be functionally pure.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the VJP rule
+    specified by ``custom_vjp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  -5.4667816
+  >>> jax.defvjp_all(f, lambda x: (np.sin(x ** 2), lambda g: (g * x,)))
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  3.0
+
+  An example with a function on two arguments, so that the VJP function must
+  return a tuple of length two:
+
+  >>> @jax.custom_transforms
+  ... def f(x, y):
+  ...   return x * y
+  ...
+  >>> jax.defvjp_all(f, lambda x, y: (x * y, lambda g: (y, x)))
+  >>> print(f(3., 4.))
+  12.0
+  >>> print(jax.grad(f, argnums=(0, 1))(3., 4.))
+  (4.0, 3.0)
+  """
+  _check_custom_transforms_type("defvjp_all", fun)
+  def custom_transforms_vjp(jax_kwargs, *jax_args, **params):
+    if jax_kwargs:
+      msg = ("defvjp_all requires the corresponding custom_transforms function "
+             "not to be called with keyword arguments.")
+      raise ValueError(msg)
+    args = map(build_tree, params['in_trees'], jax_args)
+    pytree_out, vjp_pytree = custom_vjp(*args)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    params['out_tree'].store(out_tree)  # linear_util style side effect
+    def vjp_pytree_(ct):
+      args_cts = tuple(vjp_pytree(ct))
+      if len(args_cts) != len(params['in_trees']):
+        msg = ("custom VJP function must return a tuple of length equal to the "
+               "number of positional arguments to the function being "
+               "differentiated: expected {}, got {}")
+        raise TypeError(msg.format(len(params['in_trees']), len(args_cts)))
+      return ({},) + args_cts
+    vjp, _ = pytree_fun_to_jaxtupletree_fun(lu.wrap_init(vjp_pytree_), (out_tree,))
+    return out, vjp.call_wrapped
+  ad.defvjp_all(fun.prim, custom_transforms_vjp)
+
+def defvjp(fun, *vjprules):
+  """Define VJP rules for each argument separately.
+
+  This function is a convenience wrapper around ``jax.defvjp_all`` for
+  separately defining VJP rules for each of the function's arguments. This
+  convenience wrapper does not provide a mechanism for depending on anything
+  other than the function arguments, so VJP rules defined in this way can't
+  share work with the primal computation or with each other, though those things
+  are possible using ``jax.defvjp_all`` or ``jax.custom_gradient``. See
+  ``jax.defvjp2`` for a variant that allows for defining rules that depend on
+  the output primal value. See also the ``jax.defvjp_all`` docstring.
+
+  The signature of each component VJP rule is ``lambda g, *primals: ...`` where
+  ``g`` represents the output cotangent and ``*primals`` represents all the
+  primal positional arguments.
+
+  Args:
+    fun: a custom_transforms function.
+    *vjprules: a sequence of functions or Nones specifying the VJP rule for each
+      corresponding positional argument. When an element is None, it indicates
+      that the Jacobian from the corresponding input to the output is zero.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the VJP rule
+    specified by ``*vjprules``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  >>> jax.defvjp(f, lambda g, x: 8. * g)
+  >>> print(jax.grad(f)(3.))
+  8.0
+
+  An example with a function on two arguments:
+
+  >>> @jax.custom_transforms
+  ... def f(x, y):
+  ...   return np.sin(x ** 2 + y)
+  ...
+  >>> print(f(3., 4.))
+  0.42016703
+  >>> print(jax.grad(f)(3., 4.))
+  5.4446807
+  >>> print(jax.grad(f, 1)(3., 4.))
+  0.9074468
+  >>> jax.defvjp(f, None, lambda g, x, y: g + x + y)
+  >>> print(jax.grad(f)(3., 4.))
+  0.0
+  >>> print(jax.grad(f, 1)(3., 4.))
+  8.0
+  """
+  _check_custom_transforms_type("defvjp", fun)
+  def custom_vjp(*primals):
+    ans = fun(*primals)
+    # TODO(mattjj): avoid instantiating zeros?
+    vjpfun = lambda ct: [vjp(ct, *primals) if vjp else ad_util.zeros_like_jaxval(x)
+                         for x, vjp in zip(primals, vjprules)]
+    return ans, vjpfun
+  defvjp_all(fun, custom_vjp)
+
+def defvjp2(fun, *vjprules):
+  """Define VJP rules for each argument separately.
+
+  This function is a convenience wrapper around ``jax.defvjp_all`` for
+  separately defining VJP rules for each of the function's arguments. This
+  convenience wrapper does not provide a mechanism for depending on anything
+  other than the function arguments and its primal output value.
+
+  The signature of each component VJP rule is ``lambda g, ans, *primals: ...``
+  where ``g`` represents the output cotangent, ``ans`` represents the output
+  primal, and ``*primals`` represents all the primal positional arguments.
+
+  Args:
+    fun: a custom_transforms function.
+    *vjprules: a sequence of functions or Nones specifying the VJP rule for each
+      corresponding positional argument. When an element is None, it indicates
+      that the Jacobian from the corresponding input to the output is zero.
+
+  Returns:
+    None. A side-effect is that ``fun`` is associated with the VJP rule
+    specified by ``*vjprules``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x, y):
+  ...   return np.sin(x ** 2 + y)
+  ...
+  >>> print(f(3., 4.))
+  0.42016703
+  >>> print(jax.grad(f)(3., 4.))
+  5.4446807
+  >>> print(jax.grad(f, 1)(3., 4.))
+  0.9074468
+  >>> jax.defvjp2(f, None, lambda g, ans, x, y: g + x + y + ans)
+  >>> print(jax.grad(f)(3., 4.))
+  0.0
+  >>> print(jax.grad(f, 1)(3., 4.))
+  8.420167
+  """
+  _check_custom_transforms_type("defvjp2", fun)
+  def custom_vjp(*primals):
+    ans = fun(*primals)
+    # TODO(mattjj): avoid instantiating zeros?
+    vjpfun = lambda ct: [vjp(ct, ans, *primals) if vjp else ad_util.zeros_like_jaxval(x)
+                         for x, vjp in zip(primals, vjprules)]
+    return ans, vjpfun
+  defvjp_all(fun, custom_vjp)
+
+def custom_gradient(fun):
+  """Convenience function for defining custom VJP rules (aka custom gradients).
+
+  While the canonical way to define custom VJP rules is via ``jax.defvjp_all``
+  and its convenience wrappers, the ``custom_gradient`` convenience wrapper
+  follows TensorFlow's tf.custom_gradient API.
+
+  See https://www.tensorflow.org/api_docs/python/tf/custom_gradient.
+
+  If the function to be differentiated has type signature ``a -> b``, then the
path: jax/api.py
text: @user1:
"the function to be differentiated" is kind of confusing because there's not necessarily a Python function corresponding to it (i.e. it's not `fun`). Maybe... "the primal computation"? Not sure if I'm using these terms correctly...

@author:
There's not a Python function (i.e. a Python callable) but there is a mathematical function. I agree the phrasing seems tough but I'm not sure "the primal computation" is any better than "the function to be differentiated"...

@user1:
Maybe just "the mathematical function to be differentiated"?

@author:
Done.

diff_hunk: @@ -952,25 +954,608 @@ def _valid_jaxtype(arg):
     return True
 
 
+class CustomTransformsFunction(object):
+  def __init__(self, fun, prim):
+    self.fun = fun
+    self.prim = prim
+    wraps(fun)(self)
+
+  def __repr__(self):
+    return '<jax.custom_transforms function {fun}>'.format(fun=self.__name__)
+
+  def __call__(self, *args, **kwargs):
+    jax_args, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
+    jax_kwargs, kwargs_tree = pytree_to_jaxtupletree(kwargs)
+    out_tree = lu.Store()
+    ans = self.prim.bind(jax_kwargs, *jax_args, kwargs_tree=kwargs_tree,
+                         in_trees=in_trees, out_tree=out_tree)
+    return build_tree(out_tree.val, ans)
+
 def custom_transforms(fun):
-  name = getattr(fun, '__name__', '<unnamed user primitive>')
+  """Wraps a function so that its transformation behavior can be controlled.
+
+  A primary use case of ``custom_transforms`` is defining custom VJP rules (aka
+  custom gradients) for a Python function, while still supporting other
+  transformations like ``jax.jit`` and ``jax.vmap``. Custom differentiation
+  rules can be supplied using the ``jax.defjvp`` and ``jax.defvjp`` functions.
+
+  JAX transforms Python functions by tracing which primitive operations (like
+  black-box numerical kernels) are applied to the function's input arguments to
+  produce its output, then transforming that recorded composition of primitives
+  using a table of transformation rules for each primitive. That is, for each
+  transformation (e.g. ``jax.jvp`` or ``jax.vmap``) there is a table of rules
+  that specifies how each primitive is transformed. When a new primitive is
+  added to the system, it also needs a rule for each transformation we might
+  want to apply to it.
+
+  But in some cases instead of introducing a new primitive and specifying every
+  transformation rule, we might start with a traceable/transformable Python
+  function and want to override just its VJP (i.e. give it a custom gradient).
+  More concretely, we might start with a Python function to which we can already
+  apply ``jax.jit``, but we want to change how ``jax.grad`` behaves on it. These
+  are the cases that ``custom_transforms`` handles.
+
+  The ``custom_transforms`` decorator wraps ``fun`` so that its transformation
+  behavior can be overridden, but not all transformation rules need to be
+  specified manually. For rules that aren't overridden, a default rule is used
+  for each transformation that relies on tracing ``fun``.
+
+  Args:
+    fun: a Python callable. Must be functionally pure. Its arguments and return
+      value should be arrays, scalars, or (nested) standard Python containers
+      (tuple/list/dict) thereof.
+
+  Returns:
+    A Python callable with the same input/output and transformation behavior as
+    ``fun``, but for which custom transformation rules can be supplied, e.g.
+    using ``jax.defvjp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  -5.4667816
+  >>> jax.defvjp(f, lambda g, x: g * x)
+  >>> print(jax.grad(f)(3.))
+  3.0
+  """
+  name = getattr(fun, '__name__', '<unnamed custom_transforms primitive>')
   fun_p = core.Primitive(name)
-  fun_p.def_impl(fun)
 
-  # generic transformation implementations that rely on traceability of `fun`
-  fun_p.def_abstract_eval(partial(pe.abstract_eval_fun, fun))
-  xla.translations[fun_p] = partial(xla.lower_fun, fun)
-  ad.primitive_jvps[fun_p] = partial(jvp, fun)
-  # TODO(mattjj): batching
+  def fun_impl(jax_kwargs, *jax_args, **params):
+    args = map(build_tree, params.pop('in_trees'), jax_args)
+    kwargs = build_tree(params.pop('kwargs_tree'), jax_kwargs)
+    pytree_out = fun(*args, **kwargs)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    params.pop('out_tree').store(out_tree)  # linear_util style side effect
+    assert not params
+    return out
+  fun_p.def_impl(fun_impl)
+
+  def fun_jvp(primals, tangents, **params):
+    return ad.jvp(lu.wrap_init(fun_impl, params)).call_wrapped(primals, tangents)
+  ad.primitive_jvps[fun_p] = fun_jvp
+
+  def fun_batch(batched_args, batch_dims, **params):
+    out = batching.batch(lu.wrap_init(fun_impl, params), batched_args, batch_dims, 0)
+    return out, 0
+  batching.primitive_batchers[fun_p] = fun_batch
+
+  staged_fun_p = core.Primitive('staged_' + name)
+  def fun_partial_eval(trace, *tracers, **params):
+    tracers = tuple(map(trace.instantiate_const, tracers))
+    avals = [t.aval for t in tracers]
+    pvals_in = [pe.PartialVal((a, core.unit)) for a in avals]
+    jaxpr, pval_out, consts = pe.trace_to_jaxpr(lu.wrap_init(fun_impl, params),
+                                                pvals_in, instantiate=True)
+    consts = trace.new_instantiated_const(core.pack(consts))
+    eqn = pe.JaxprEqn((consts,) + tracers, None, staged_fun_p, (), False, False,
+                      dict(params, jaxpr=jaxpr))
+    return pe.JaxprTracer(trace, pval_out, eqn)
+  pe.custom_partial_eval_rules[fun_p] = fun_partial_eval
+
+  def staged_fun_translation(c, xla_consts, *xla_args, **params):
+    consts_shapes = tuple(c.GetShape(xla_consts).tuple_shapes())
+    xla_consts = tuple(xla.xla_destructure(c, xla_consts))
+    arg_shapes = map(c.GetShape, xla_args)
+    built_c = xla.jaxpr_computation(params['jaxpr'], (), consts_shapes, *arg_shapes)
+    return c.Call(built_c, xla_consts + xla_args)
+  xla.translations[staged_fun_p] = staged_fun_translation
+
+  return CustomTransformsFunction(fun, fun_p)
+
+def _check_custom_transforms_type(name, fun):
+  if type(fun) is not CustomTransformsFunction:
+    msg = ("{} requires a custom_transforms function as its first argument, "
+          "but got type {}.")
+    raise TypeError(msg.format(name, type(fun)))
+
+def defjvp_all(fun, custom_jvp):
+  """Define a custom JVP rule for a ``custom_transforms`` function.
+
+  If ``fun`` represents a function with signature ``a -> b``, then
+  ``custom_jvp`` represents a function with signature ``a -> T a -> (b, Tb)``,
+  where we use ``T x`` to represent a tangent type for the type ``x``.
+
+  In more detail, ``custom_jvp`` must two arguments, both tuples of length equal
+  to the number of positional arguments to ``fun``. The first argument to
+  ``custom_jvp`` represents the input primal values, and the second represents
+  the input tangent values. ``custom_jvp`` must return a pair where the first
+  element represents the output primal value and the second element represents
+  the output tangent value.
+
+  Defining a custom JVP rule also affects the default VJP rule, which is derived
+  from the JVP rule automatically via transposition.
 
-  @wraps(fun)
-  def traceable(*args, **kwargs):
-    # TODO(mattjj): pytrees to jaxtupletrees
-    return fun_p.bind(*args, **kwargs)
-  traceable.primitive = fun_p
+  Args:
+    fun: a custom_transforms function.
+    custom_jvp: a Python callable specifying the JVP rule, taking two tuples as
+      arguments specifying the input primal values and tangent values,
+      respectively. The tuple elements can be arrays, scalars, or (nested)
+      standard Python containers (tuple/list/dict) thereof. Must be functionally
+      pure.
path: jax/api.py
text: @user1:
Can the output be a pytree? Also briefly mention what the output is here.

@author:
Done.

diff_hunk: @@ -952,25 +954,608 @@ def _valid_jaxtype(arg):
     return True
 
 
+class CustomTransformsFunction(object):
+  def __init__(self, fun, prim):
+    self.fun = fun
+    self.prim = prim
+    wraps(fun)(self)
+
+  def __repr__(self):
+    return '<jax.custom_transforms function {fun}>'.format(fun=self.__name__)
+
+  def __call__(self, *args, **kwargs):
+    jax_args, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
+    jax_kwargs, kwargs_tree = pytree_to_jaxtupletree(kwargs)
+    out_tree = lu.Store()
+    ans = self.prim.bind(jax_kwargs, *jax_args, kwargs_tree=kwargs_tree,
+                         in_trees=in_trees, out_tree=out_tree)
+    return build_tree(out_tree.val, ans)
+
 def custom_transforms(fun):
-  name = getattr(fun, '__name__', '<unnamed user primitive>')
+  """Wraps a function so that its transformation behavior can be controlled.
+
+  A primary use case of ``custom_transforms`` is defining custom VJP rules (aka
+  custom gradients) for a Python function, while still supporting other
+  transformations like ``jax.jit`` and ``jax.vmap``. Custom differentiation
+  rules can be supplied using the ``jax.defjvp`` and ``jax.defvjp`` functions.
+
+  JAX transforms Python functions by tracing which primitive operations (like
+  black-box numerical kernels) are applied to the function's input arguments to
+  produce its output, then transforming that recorded composition of primitives
+  using a table of transformation rules for each primitive. That is, for each
+  transformation (e.g. ``jax.jvp`` or ``jax.vmap``) there is a table of rules
+  that specifies how each primitive is transformed. When a new primitive is
+  added to the system, it also needs a rule for each transformation we might
+  want to apply to it.
+
+  But in some cases instead of introducing a new primitive and specifying every
+  transformation rule, we might start with a traceable/transformable Python
+  function and want to override just its VJP (i.e. give it a custom gradient).
+  More concretely, we might start with a Python function to which we can already
+  apply ``jax.jit``, but we want to change how ``jax.grad`` behaves on it. These
+  are the cases that ``custom_transforms`` handles.
+
+  The ``custom_transforms`` decorator wraps ``fun`` so that its transformation
+  behavior can be overridden, but not all transformation rules need to be
+  specified manually. For rules that aren't overridden, a default rule is used
+  for each transformation that relies on tracing ``fun``.
+
+  Args:
+    fun: a Python callable. Must be functionally pure. Its arguments and return
+      value should be arrays, scalars, or (nested) standard Python containers
+      (tuple/list/dict) thereof.
+
+  Returns:
+    A Python callable with the same input/output and transformation behavior as
+    ``fun``, but for which custom transformation rules can be supplied, e.g.
+    using ``jax.defvjp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  -5.4667816
+  >>> jax.defvjp(f, lambda g, x: g * x)
+  >>> print(jax.grad(f)(3.))
+  3.0
+  """
+  name = getattr(fun, '__name__', '<unnamed custom_transforms primitive>')
   fun_p = core.Primitive(name)
-  fun_p.def_impl(fun)
 
-  # generic transformation implementations that rely on traceability of `fun`
-  fun_p.def_abstract_eval(partial(pe.abstract_eval_fun, fun))
-  xla.translations[fun_p] = partial(xla.lower_fun, fun)
-  ad.primitive_jvps[fun_p] = partial(jvp, fun)
-  # TODO(mattjj): batching
+  def fun_impl(jax_kwargs, *jax_args, **params):
+    args = map(build_tree, params.pop('in_trees'), jax_args)
+    kwargs = build_tree(params.pop('kwargs_tree'), jax_kwargs)
+    pytree_out = fun(*args, **kwargs)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    params.pop('out_tree').store(out_tree)  # linear_util style side effect
+    assert not params
+    return out
+  fun_p.def_impl(fun_impl)
+
+  def fun_jvp(primals, tangents, **params):
+    return ad.jvp(lu.wrap_init(fun_impl, params)).call_wrapped(primals, tangents)
+  ad.primitive_jvps[fun_p] = fun_jvp
+
+  def fun_batch(batched_args, batch_dims, **params):
+    out = batching.batch(lu.wrap_init(fun_impl, params), batched_args, batch_dims, 0)
+    return out, 0
+  batching.primitive_batchers[fun_p] = fun_batch
+
+  staged_fun_p = core.Primitive('staged_' + name)
+  def fun_partial_eval(trace, *tracers, **params):
+    tracers = tuple(map(trace.instantiate_const, tracers))
+    avals = [t.aval for t in tracers]
+    pvals_in = [pe.PartialVal((a, core.unit)) for a in avals]
+    jaxpr, pval_out, consts = pe.trace_to_jaxpr(lu.wrap_init(fun_impl, params),
+                                                pvals_in, instantiate=True)
+    consts = trace.new_instantiated_const(core.pack(consts))
+    eqn = pe.JaxprEqn((consts,) + tracers, None, staged_fun_p, (), False, False,
+                      dict(params, jaxpr=jaxpr))
+    return pe.JaxprTracer(trace, pval_out, eqn)
+  pe.custom_partial_eval_rules[fun_p] = fun_partial_eval
+
+  def staged_fun_translation(c, xla_consts, *xla_args, **params):
+    consts_shapes = tuple(c.GetShape(xla_consts).tuple_shapes())
+    xla_consts = tuple(xla.xla_destructure(c, xla_consts))
+    arg_shapes = map(c.GetShape, xla_args)
+    built_c = xla.jaxpr_computation(params['jaxpr'], (), consts_shapes, *arg_shapes)
+    return c.Call(built_c, xla_consts + xla_args)
+  xla.translations[staged_fun_p] = staged_fun_translation
+
+  return CustomTransformsFunction(fun, fun_p)
+
+def _check_custom_transforms_type(name, fun):
+  if type(fun) is not CustomTransformsFunction:
+    msg = ("{} requires a custom_transforms function as its first argument, "
+          "but got type {}.")
+    raise TypeError(msg.format(name, type(fun)))
+
+def defjvp_all(fun, custom_jvp):
+  """Define a custom JVP rule for a ``custom_transforms`` function.
+
+  If ``fun`` represents a function with signature ``a -> b``, then
+  ``custom_jvp`` represents a function with signature ``a -> T a -> (b, Tb)``,
path: jax/api.py
text: @user1:
I still don't think the Haskell notation will be comprehensible to most users, but maybe it's part of learning jax :) At least do `T a` + `T b` (or `Ta` + `Tb`) to make the tangent notation consistent.

Edit: I see you explain it in words below, PHEW

@author:
Done.

diff_hunk: @@ -952,25 +954,608 @@ def _valid_jaxtype(arg):
     return True
 
 
+class CustomTransformsFunction(object):
+  def __init__(self, fun, prim):
+    self.fun = fun
+    self.prim = prim
+    wraps(fun)(self)
+
+  def __repr__(self):
+    return '<jax.custom_transforms function {fun}>'.format(fun=self.__name__)
+
+  def __call__(self, *args, **kwargs):
+    jax_args, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
+    jax_kwargs, kwargs_tree = pytree_to_jaxtupletree(kwargs)
+    out_tree = lu.Store()
+    ans = self.prim.bind(jax_kwargs, *jax_args, kwargs_tree=kwargs_tree,
+                         in_trees=in_trees, out_tree=out_tree)
+    return build_tree(out_tree.val, ans)
+
 def custom_transforms(fun):
-  name = getattr(fun, '__name__', '<unnamed user primitive>')
+  """Wraps a function so that its transformation behavior can be controlled.
+
+  A primary use case of ``custom_transforms`` is defining custom VJP rules (aka
+  custom gradients) for a Python function, while still supporting other
+  transformations like ``jax.jit`` and ``jax.vmap``. Custom differentiation
+  rules can be supplied using the ``jax.defjvp`` and ``jax.defvjp`` functions.
+
+  JAX transforms Python functions by tracing which primitive operations (like
+  black-box numerical kernels) are applied to the function's input arguments to
+  produce its output, then transforming that recorded composition of primitives
+  using a table of transformation rules for each primitive. That is, for each
+  transformation (e.g. ``jax.jvp`` or ``jax.vmap``) there is a table of rules
+  that specifies how each primitive is transformed. When a new primitive is
+  added to the system, it also needs a rule for each transformation we might
+  want to apply to it.
+
+  But in some cases instead of introducing a new primitive and specifying every
+  transformation rule, we might start with a traceable/transformable Python
+  function and want to override just its VJP (i.e. give it a custom gradient).
+  More concretely, we might start with a Python function to which we can already
+  apply ``jax.jit``, but we want to change how ``jax.grad`` behaves on it. These
+  are the cases that ``custom_transforms`` handles.
+
+  The ``custom_transforms`` decorator wraps ``fun`` so that its transformation
+  behavior can be overridden, but not all transformation rules need to be
+  specified manually. For rules that aren't overridden, a default rule is used
+  for each transformation that relies on tracing ``fun``.
+
+  Args:
+    fun: a Python callable. Must be functionally pure. Its arguments and return
+      value should be arrays, scalars, or (nested) standard Python containers
+      (tuple/list/dict) thereof.
+
+  Returns:
+    A Python callable with the same input/output and transformation behavior as
+    ``fun``, but for which custom transformation rules can be supplied, e.g.
+    using ``jax.defvjp``.
+
+  For example:
+
+  >>> @jax.custom_transforms
+  ... def f(x):
+  ...   return np.sin(x ** 2)
+  ...
+  >>> print(f(3.))
+  0.4121185
+  >>> print(jax.grad(f)(3.))
+  -5.4667816
+  >>> jax.defvjp(f, lambda g, x: g * x)
+  >>> print(jax.grad(f)(3.))
+  3.0
+  """
+  name = getattr(fun, '__name__', '<unnamed custom_transforms primitive>')
   fun_p = core.Primitive(name)
-  fun_p.def_impl(fun)
 
-  # generic transformation implementations that rely on traceability of `fun`
-  fun_p.def_abstract_eval(partial(pe.abstract_eval_fun, fun))
-  xla.translations[fun_p] = partial(xla.lower_fun, fun)
-  ad.primitive_jvps[fun_p] = partial(jvp, fun)
-  # TODO(mattjj): batching
+  def fun_impl(jax_kwargs, *jax_args, **params):
+    args = map(build_tree, params.pop('in_trees'), jax_args)
+    kwargs = build_tree(params.pop('kwargs_tree'), jax_kwargs)
+    pytree_out = fun(*args, **kwargs)
+    out, out_tree = pytree_to_jaxtupletree(pytree_out)
+    params.pop('out_tree').store(out_tree)  # linear_util style side effect
+    assert not params
+    return out
+  fun_p.def_impl(fun_impl)
+
+  def fun_jvp(primals, tangents, **params):
+    return ad.jvp(lu.wrap_init(fun_impl, params)).call_wrapped(primals, tangents)
+  ad.primitive_jvps[fun_p] = fun_jvp
+
+  def fun_batch(batched_args, batch_dims, **params):
+    out = batching.batch(lu.wrap_init(fun_impl, params), batched_args, batch_dims, 0)
+    return out, 0
+  batching.primitive_batchers[fun_p] = fun_batch
+
+  staged_fun_p = core.Primitive('staged_' + name)
+  def fun_partial_eval(trace, *tracers, **params):
+    tracers = tuple(map(trace.instantiate_const, tracers))
+    avals = [t.aval for t in tracers]
+    pvals_in = [pe.PartialVal((a, core.unit)) for a in avals]
+    jaxpr, pval_out, consts = pe.trace_to_jaxpr(lu.wrap_init(fun_impl, params),
+                                                pvals_in, instantiate=True)
+    consts = trace.new_instantiated_const(core.pack(consts))
+    eqn = pe.JaxprEqn((consts,) + tracers, None, staged_fun_p, (), False, False,
+                      dict(params, jaxpr=jaxpr))
+    return pe.JaxprTracer(trace, pval_out, eqn)
+  pe.custom_partial_eval_rules[fun_p] = fun_partial_eval
+
+  def staged_fun_translation(c, xla_consts, *xla_args, **params):
+    consts_shapes = tuple(c.GetShape(xla_consts).tuple_shapes())
+    xla_consts = tuple(xla.xla_destructure(c, xla_consts))
+    arg_shapes = map(c.GetShape, xla_args)
+    built_c = xla.jaxpr_computation(params['jaxpr'], (), consts_shapes, *arg_shapes)
+    return c.Call(built_c, xla_consts + xla_args)
+  xla.translations[staged_fun_p] = staged_fun_translation
+
+  return CustomTransformsFunction(fun, fun_p)
+
+def _check_custom_transforms_type(name, fun):
+  if type(fun) is not CustomTransformsFunction:
+    msg = ("{} requires a custom_transforms function as its first argument, "
+          "but got type {}.")
+    raise TypeError(msg.format(name, type(fun)))
+
+def defjvp_all(fun, custom_jvp):
+  """Define a custom JVP rule for a ``custom_transforms`` function.
+
+  If ``fun`` represents a function with signature ``a -> b``, then
+  ``custom_jvp`` represents a function with signature ``a -> T a -> (b, Tb)``,
+  where we use ``T x`` to represent a tangent type for the type ``x``.
+
+  In more detail, ``custom_jvp`` must two arguments, both tuples of length equal
path: jax/api.py
text: @user1:
must take two

@author:
Done.

diff_hunk: @@ -952,25 +954,608 @@ def _valid_jaxtype(arg):
     return True
 
 
+class CustomTransformsFunction(object):
+  def __init__(self, fun, prim):
+    self.fun = fun
+    self.prim = prim
+    wraps(fun)(self)
+
+  def __repr__(self):
+    return '<jax.custom_transforms function {fun}>'.format(fun=self.__name__)
+
+  def __call__(self, *args, **kwargs):
+    jax_args, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
+    jax_kwargs, kwargs_tree = pytree_to_jaxtupletree(kwargs)
+    out_tree = lu.Store()
+    ans = self.prim.bind(jax_kwargs, *jax_args, kwargs_tree=kwargs_tree,
+                         in_trees=in_trees, out_tree=out_tree)
+    return build_tree(out_tree.val, ans)
+
 def custom_transforms(fun):
-  name = getattr(fun, '__name__', '<unnamed user primitive>')
+  """Wraps a function so that its transformation behavior can be controlled.
+
+  A primary use case of ``custom_transforms`` is defining custom VJP rules (aka
+  custom gradients) for a Python function, while still supporting other
+  transformations like ``jax.jit`` and ``jax.vmap``. Custom differentiation
+  rules can be supplied using the ``jax.defjvp`` and ``jax.defvjp`` functions.
+
+  JAX transforms Python functions by tracing which primitive operations (like
+  black-box numerical kernels) are applied to the function's input arguments to
+  produce its output, then transforming that recorded composition of primitives
+  using a table of transformation rules for each primitive. That is, for each
+  transformation (e.g. ``jax.jvp`` or ``jax.vmap``) there is a table of rules
+  that specifies how each primitive is transformed. When a new primitive is
+  added to the system, it also needs a rule for each transformation we might
+  want to apply to it.
+
+  But in some cases instead of introducing a new primitive and specifying every
+  transformation rule, we might start with a traceable/transformable Python
+  function and want to override just its VJP (i.e. give it a custom gradient).
+  More concretely, we might start with a Python function to which we can already
+  apply ``jax.jit``, but we want to change how ``jax.grad`` behaves on it. These
+  are the cases that ``custom_transforms`` handles.
+
+  The ``custom_transforms`` decorator wraps ``fun`` so that its transformation
+  behavior can be overridden, but not all transformation rules need to be
+  specified manually. For rules that aren't overridden, a default rule is used
+  for each transformation that relies on tracing ``fun``.
path: jax/api.py
text: @user1:
Maybe replace last sentence with "The default behavior is retained for any non-overridden rules."

@author:
Done.