Repository: google/jax

Issue:
Deferred pushforward evaluation in jvp
Hi. I'm sorry if this question is due to a basic misunderstanding of how the AD in JAX works. I tried to work through the code but found it pretty tough to follow.

Basically, for vjps we get back a function to evaluate the pullback of a cotangent vector without re-evaluating the original function (and, I imagine, recomputing intermediate values). I'm wondering if there is something similar for jvps? I'd like to get back a function to evaluate the pushforward of a tangent vector at a specific point, ideally making use of intermediate values computed when evaluating the function at the chosen point. My goal is to evaluate the pushforward at the same point but for multiple tangent vectors but where these tangent vectors are not known a priori as they are in jacfwd, so vmap isn't an option. Is there a better way to do this than simply calling jvp every time?

I see that there is a linearize function that looks like it gives me the pushforward function I want but it's not documented so I'm wondering if it's meant for internal use?

Thank you in advance

Deferred pushforward evaluation in jvp
Hi. I'm sorry if this question is due to a basic misunderstanding of how the AD in JAX works. I tried to work through the code but found it pretty tough to follow.

Basically, for vjps we get back a function to evaluate the pullback of a cotangent vector without re-evaluating the original function (and, I imagine, recomputing intermediate values). I'm wondering if there is something similar for jvps? I'd like to get back a function to evaluate the pushforward of a tangent vector at a specific point, ideally making use of intermediate values computed when evaluating the function at the chosen point. My goal is to evaluate the pushforward at the same point but for multiple tangent vectors but where these tangent vectors are not known a priori as they are in jacfwd, so vmap isn't an option. Is there a better way to do this than simply calling jvp every time?

I see that there is a linearize function that looks like it gives me the pushforward function I want but it's not documented so I'm wondering if it's meant for internal use?

Thank you in advance



Discussion:
Thanks for asking this!

You're right, that's exactly what `linearize` is for. It composes forward-mode autodiff with partial evaluation, so that all the linearization points are stored (which costs memory, but means you don't have to re-do FLOPs for future JVP evaluations).

Everything in api.py (where the name doesn't start with an underscore) is public and meant to be used. We just haven't gotten to documenting some things.

@sschoenholz anything to add?

Maybe we can use this issue to add a docstring to `linearize`.
Here's some quick example usage:

```python
from __future__ import print_function

import jax.numpy as np
from jax import jvp, linearize

def f(x):
  return 3. * np.sin(x) + np.cos(x / 2.)

x = 2.
t1 = 1.
t2 = 3.

print(jvp(f, (x,), (t1,)))
print(jvp(f, (x,), (t2,)))  # nothing saved from first evaluation


y, f_jvp = linearize(f, x)
print(y)
print(f_jvp(t1))  # all the linearization work is already done!
print(f_jvp(t2))
```
Thanks for asking this!

You're right, that's exactly what `linearize` is for. It composes forward-mode autodiff with partial evaluation, so that all the linearization points are stored (which costs memory, but means you don't have to re-do FLOPs for future JVP evaluations).

Everything in api.py (where the name doesn't start with an underscore) is public and meant to be used. We just haven't gotten to documenting some things.

@sschoenholz anything to add?

Maybe we can use this issue to add a docstring to `linearize`.
Here's some quick example usage:

```python
from __future__ import print_function

import jax.numpy as np
from jax import jvp, linearize

def f(x):
  return 3. * np.sin(x) + np.cos(x / 2.)

x = 2.
t1 = 1.
t2 = 3.

print(jvp(f, (x,), (t1,)))
print(jvp(f, (x,), (t2,)))  # nothing saved from first evaluation


y, f_jvp = linearize(f, x)
print(y)
print(f_jvp(t1))  # all the linearization work is already done!
print(f_jvp(t2))
```

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -523,10 +523,62 @@ def trim_arg(primal, tangent):
   out_primal, out_tangent = ad.jvp(jaxtree_fun).call_wrapped(ps_flat, ts_flat)
   return (build_tree(out_tree(), out_primal), build_tree(out_tree(), out_tangent))
 
-def linearize(traceable, *primals):
-  fun = lu.wrap_init(traceable)
+def linearize(fun, *primals):
+  """Produce a linear approximation to `fun` using `jvp` and partial evaluation.
+
+  Args:
+    fun: Function to be differentiated. Its arguments should be arrays, scalars,
+      or standard Python containers of arrays or scalars. It should return an
+      array, scalar, or standard python container of arrays or scalars.
+    primals: The primal values at which the Jacobian of `fun` should be
+      evaluated. Should be a tuple of arrays, scalar, or standard Python
+      container thereof. The length of the tuple is equal to the number of
+      positional parameters of `fun`.
+
+  Returns:
+    A pair where the first element is the value of `f(*primals)` and the second
+    element is a function that evaluates the (forward-mode) Jacobian-vector
+    product of `fun` evaluated at `primals` without re-doing the linearization
+    work.
+
+  In terms of values computed, `linearize` behaves much like a curried `jvp`::
+    y, out_tangent = jax.jvp(f, (x,), (in_tangent,))
+
+    y, f_jvp = jax.linearize(f, x)
+    out_tangent = f_jvp(in_tangent)
+
+  However, the difference is that `linearize` uses partial evaluation so that
+  the function `f` is not re-linearized on calls to `f_jvp`. In general that
+  means the memory usage scales with the size of the computation, much like in
+  reverse-mode. (Indeed, `linearize` has a similar signature to `vjp`!)
+
+  This function is mainly useful if you want to apply `f_jvp` multiple times,
+  i.e. to evaluate a pushforward for many different input tangent vectors at the
+  same linearization point. Moreover if all the input tangent vectors are known
+  at once, it can be more efficient to vectorize using `vmap`, as in::
+    pushfwd = partial(jvp, f, (x,))
+    y, out_tangents = vmap(pushfwd, out_axes=(None, 0))((in_tangents,))
+  By using `vmap` and `jvp` together like this we avoid the stored-linearization
+  memory cost that scales with the depth of the computation, which is incurred
+  by both `linearize` and `vjp`.
+
+  Here's a more complete example of using `linearize`:
+
+  >>> def f(x): return 3. * np.sin(x) + np.cos(x / 2.)
+  ...
+  >>> jax.jvp(f, (2.,), (3.,))
+  (array(3.2681944, dtype=float32), array(-5.007528, dtype=float32))
+  >>> y, f_jvp = jax.linearize(f, 2.)
+  >>> y
+  array(3.2681944, dtype=float32)
+  >>> f_jvp(3.)
+  array(-5.007528, dtype=float32)
+  >>> f_jvp(4.)
+  array(-5.007528, dtype=float32)
+  """
+  f = lu.wrap_init(fun)
   primals_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, primals))
-  jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(fun, in_trees)
+  jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(f, in_trees)
   out_primal, out_pval, jaxpr, consts = ad.linearize(jaxtree_fun, *primals_flat)
   out_tree = out_tree()
   out_primal_py = build_tree(out_tree, out_primal)


Test patch: