Repository: google/jax

Issue:
FFT Hessian broken
Been writing some FFT code lately so hitting some bumps in the road--it looks like Hessians are broken for FFTs, for now. It's blocking my research project, so I'd be happy to work on this, with some guidance for what to add (not sure if it's adding a batching rule or something tricker related to #1021). But, it would probably take less time for someone already familiar with the codebase.

Reproducing example:

```python
import jax
import jax.numpy as np

fft_hess = jax.hessian(np.fft.fftn)
fft_hess(np.array([1.0, 2.0]))  # Errors
```

This produces the following stack trace:

```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/batching.py in get_primitive_batcher(p)
    231   try:
--> 232     return primitive_batchers[p]
    233   except KeyError:

KeyError: fft

During handling of the above exception, another exception occurred:

NotImplementedError                       Traceback (most recent call last)
 in 
----> 1 fft_hess(np.array([1.0, 2.0]))

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/api.py in jacfun(*args, **kwargs)
    342     holomorphic or tree_map(_check_real_input_jacfwd, dyn_args)
    343     pushfwd = partial(jvp, f_partial, dyn_args)
--> 344     y, jac = vmap(pushfwd, out_axes=(None, -1))(_std_basis(dyn_args))
    345     example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args
    346     return tree_map(partial(_unravel_array_into_pytree, example_args, -1), jac)

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/api.py in batched_fun(*args, **kwargs)
    491     in_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
    492     jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(f, in_trees)
--> 493     out_flat = batching.batch(jaxtree_fun, in_flat, in_axes_, out_axes)
    494     return build_tree(out_tree(), out_flat)
    495 

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/batching.py in batch(fun, in_vals, in_dims, out_dim_dst)
     43   elif len(sizes) == 1:
     44     sz = sizes.pop()
---> 45     return batch_transform(fun, sz, in_dims, out_dim_dst).call_wrapped(in_vals)
     46   else:
     47     raise TypeError("got inconsistent map dimension sizes: {}".format(sizes))

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    145 
    146     del gen
--> 147     ans = self.f(*args, **dict(self.params, **kwargs))
    148     del args
    149     while stack:

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/api.py in jvp(fun, primals, tangents)
    761   ps_flat, ts_flat, in_trees = unzip3(map(trim_arg, primals, tangents))
    762   jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(fun, in_trees)
--> 763   out_primal, out_tangent = ad.jvp(jaxtree_fun).call_wrapped(ps_flat, ts_flat)
    764   return (build_tree(out_tree(), out_primal), build_tree(out_tree(), out_tangent))
    765 

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    145 
    146     del gen
--> 147     ans = self.f(*args, **dict(self.params, **kwargs))
    148     del args
    149     while stack:

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/api.py in jacfun(*args, **kwargs)
    383     f = lu.wrap_init(fun, kwargs)
    384     f_partial, dyn_args = _argnums_partial(f, argnums, args)
--> 385     y, pullback = vjp(f_partial, *dyn_args)
    386     holomorphic or tree_map(_check_real_output_jacrev, y)
    387     jac = vmap(pullback)(_std_basis(y))

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/api.py in vjp(fun, *primals, **kwargs)
    896   jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(fun, in_trees)
    897   if not has_aux:
--> 898     out_primal, out_vjp = ad.vjp(jaxtree_fun, primals_flat)
    899   else:
    900     out_primal, out_vjp, aux = ad.vjp(jaxtree_fun, primals_flat, has_aux=True)

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux)
    103 def vjp(traceable, primals, has_aux=False):
    104   if not has_aux:
--> 105     out_primal, pval, jaxpr, consts = linearize(traceable, *primals)
    106   else:
    107     out_primal, pval, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
     92   in_pvals = (pe.PartialVal((None, pack(primals))),
     93               pe.PartialVal((core.AbstractTuple(tangent_avals), core.unit)))
---> 94   jaxpr, out_pval, consts = pe.trace_to_jaxpr(jvpfun, in_pvals)
     95   pval_primal, pval_tangent = unpair_pval(out_pval)
     96   aval_primal, const_primal = pval_primal

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, **kwargs)
    399   with new_master(JaxprTrace) as master:
    400     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 401     jaxpr, (out_pval, consts, env) = fun.call_wrapped(pvals)
    402     assert not env
    403     del master

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    145 
    146     del gen
--> 147     ans = self.f(*args, **dict(self.params, **kwargs))
    148     del args
    149     while stack:

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/numpy/fft.py in fftn(a, s, axes, norm)
     68     s = [a.shape[axis] for axis in axes]
     69   a = _promote_to_complex(a)
---> 70   return lax.fft(a, xla_client.FftType.FFT, s)
     71 
     72 for func in get_module_functions(onp.fft):

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/lax/lax_fft.py in fft(x, fft_type, fft_lengths)
     28   else:
     29     fft_lengths = tuple(fft_lengths)
---> 30   return fft_p.bind(x, fft_type=fft_type, fft_lengths=fft_lengths)
     31 
     32 def fft_impl(x, fft_type, fft_lengths):

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/core.py in bind(self, *args, **kwargs)
    145 
    146     tracers = map(top_trace.full_raise, args)
--> 147     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    148     return full_lower(out_tracer)
    149 

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/ad.py in process_primitive(self, primitive, tracers, params)
    250           "Forward-mode differentiation rule for '{}' not implemented"
    251           .format(primitive))
--> 252     primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
    253     return JVPTracer(self, primal_out, tangent_out)
    254 

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/ad.py in linear_jvp(primitive, primals, tangents, **params)
    356 
    357 def linear_jvp(primitive, primals, tangents, **params):
--> 358   val_out = primitive.bind(*primals, **params)
    359   if all(tangent is zero for tangent in tangents):
    360     return val_out, zero

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/core.py in bind(self, *args, **kwargs)
    145 
    146     tracers = map(top_trace.full_raise, args)
--> 147     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    148     return full_lower(out_tracer)
    149 

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/ad.py in process_primitive(self, primitive, tracers, params)
    250           "Forward-mode differentiation rule for '{}' not implemented"
    251           .format(primitive))
--> 252     primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
    253     return JVPTracer(self, primal_out, tangent_out)
    254 

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/ad.py in linear_jvp(primitive, primals, tangents, **params)
    361   else:
    362     tangents = map(instantiate_zeros, primals, tangents)
--> 363     return val_out, primitive.bind(*tangents, **params)
    364 
    365 def linear_transpose(transpose_rule, cotangent, *args, **kwargs):

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/core.py in bind(self, *args, **kwargs)
    145 
    146     tracers = map(top_trace.full_raise, args)
--> 147     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    148     return full_lower(out_tracer)
    149 

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/batching.py in process_primitive(self, primitive, tracers, params)
    121     else:
    122       # TODO(mattjj,phawkins): if no rule implemented, could vmap-via-map here
--> 123       batched_primitive = get_primitive_batcher(primitive)
    124       val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
    125       return BatchTracer(self, val_out, dim_out)

~/.pyenv/versions/graf/lib/python3.7/site-packages/jax/interpreters/batching.py in get_primitive_batcher(p)
    233   except KeyError:
    234     raise NotImplementedError(
--> 235         "Batching rule for '{}' not implemented".format(p))
    236 
    237 def defvectorized(prim):

NotImplementedError: Batching rule for 'fft' not implemented
```


Discussion:
Try adding this to `lax_fft.py` (lightly tested):

```
def fft_batching_rule(batched_args, batch_dims, fft_type, fft_lengths):
    x, = batched_args
    bd, = batch_dims
    x = batching.bdim_at_front(x, bd)
    return fft(x, fft_type, fft_lengths), 0
batching.primitive_batchers[fft_p] = fft_batching_rule
```

This seems to work--many thanks for the quick response! I'll comment if any issues arise.

(feel free to close--unless you want to leave the issue open for tracking)

further edit: The JIT-ed code is ridiculously fast, put a huge smile on my face. Thanks for making a tool that makes users happy! :) 
Great! One caution about microbenchmarking JAX: make sure you read this first so you are timing the right things: https://jax.readthedocs.io/en/latest/async_dispatch.html


Resolving patch:
diff --git a/jax/lax/lax_fft.py b/jax/lax/lax_fft.py
--- a/jax/lax/lax_fft.py
+++ b/jax/lax/lax_fft.py
@@ -20,11 +20,15 @@
 from jax.core import Primitive
 from jax.interpreters import xla
 from ..interpreters import ad
+from ..interpreters import batching
 
 
 def fft(x, fft_type, fft_lengths=None):
   if fft_lengths is None:
     fft_lengths = x.shape
+  elif len(fft_lengths) == 0:
+    # XLA FFT doesn't support 0-rank.
+    return x
   else:
     fft_lengths = tuple(fft_lengths)
   return fft_p.bind(x, fft_type=fft_type, fft_lengths=fft_lengths)
@@ -41,8 +45,15 @@ def fft_translation_rule(c, x, fft_type, fft_lengths):
 def fft_transpose_rule(t, fft_type, fft_lengths):
   return fft(t, fft_type, fft_lengths),
 
+def fft_batching_rule(batched_args, batch_dims, fft_type, fft_lengths):
+  x, = batched_args
+  bd, = batch_dims
+  x = batching.bdim_at_front(x, bd)
+  return fft(x, fft_type, fft_lengths), 0
+
 fft_p = Primitive('fft')
 fft_p.def_impl(fft_impl)
 fft_p.def_abstract_eval(fft_abstract_eval)
 xla.translations[fft_p] = fft_translation_rule
 ad.deflinear(fft_p, fft_transpose_rule)
+batching.primitive_batchers[fft_p] = fft_batching_rule


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2645,6 +2645,21 @@ def fun(operand):
       for bdims in all_bdims(shape):
         self._CheckBatching(fun, 3, bdims, (shape,), dtype, rng)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_shape={}_bdims={}_fft_ndims={}"
+       .format(shape, bdims, fft_ndims),
+       "shape": shape, "bdims": bdims, "fft_ndims": fft_ndims, "rng": rng}
+      for shape in [(5,), (3, 4, 5), (2, 3, 4, 5)]
+      for bdims in all_bdims(shape)
+      for fft_ndims in range(0, min(3, len(shape)) + 1)
+      for rng in [jtu.rand_default()]))
+  def testFft(self, fft_ndims, shape, bdims, rng):
+    ndims = len(shape)
+    axes = range(ndims - fft_ndims, ndims)
+    fft_lengths = [shape[axis] for axis in axes]
+    op = lambda x: lax.fft(x, xla_bridge.xla_client.FftType.FFT, fft_lengths)
+    self._CheckBatching(op, 5, bdims, [shape], onp.complex64, rng)
+
   # TODO Concatenate
   # TODO Reverse
   # TODO DynamicSlice