Repository: google/jax

Issue:
How to use jax.hessian?
I'm using this snippet from the README

```
from jax import jit, jacfwd, jacrev
def hessian(fun):
  return jit(jacfwd(jacrev(fun)))
```

combined with the "getting started with pytorch data loaders" colab. How do I compute and use the hessian of this neural network? So far I have tried:

1. Naive: `hessian_loss = hessian(loss)(params, x, y)` (`TypeError: jacfwd() takes 2 positional arguments but 4 were given`)
2. Naive2: `hessian_loss = hessian(loss)((params, x, y))` (takes a long time with a small network and then returns `AttributeError: 'PyLeaf' object has no attribute 'node_type'`)
3. FromTheTests: `hessian_loss = hessian(lambda params_: loss(params_, x, y)(params)` (`AttributeError: 'PyLeaf' object has no attribute 'node_type'`)

All I really need is to compute the eigenvalues and eigenvectors of this Hessian.

Note: `repr(params) = [(DeviceArray{float32[10,784]}, DeviceArray{float32[10]})]`

Some ideas that I have:
1. Should I re-write the code to use only flat arrays?


Discussion:
This is a great question, and one that we've been discussing!

One problem is that JAX gave you very unhelpful error messages for this use case. We need to improve those. Ultimately, we need the error messages to communicate that the `hessian` function, as with `jacfwd` and `jacrev`, only apply to array-input array-output functions (of one argument). In particular, they don't work for tuple/list/dict inputs, or with respect to multiple arguments, for the same reason in both cases.

Here's the heart of the issue: given the example `params` here (a one-element list of a pair of arrays), how would you want the value of `hessian(lambda params: loss(params, x, y))(params)` to be stored? More generally, given a tuple/list/dict argument, how should we represent the Hessian?

For arrays, there's a clear answer because it's easy to reason about adding axes. If a function `fun` takes arrays of shape `(in_1, in_2, ..., in_n)` to arrays of shape `(out_1, out_2, ..., out_m)`, then it's reasonable for `hessian(fun)` to be a function that takes an array of shape `(in_1, in_2, ..., in_n)` to an array of shape `(out_1, out_2, ..., out_m, in_1, in_2, ..., in_n, in_1, in_2, ..., in_n)`, though other conventions could be reasonable too. (As I wrote this, I got a sense of [deja vu](https://github.com/HIPS/autograd/issues/363#issuecomment-366559091)...)

But if `fun` is, say, a function that takes a tuple of scalars to a scalar, then what should `hessian(fun)` return? Some kind of nested tuple structure? How do we organize the nesting?

I really mean that as a question! Do you have a clear sense of what would make sense from your perspective? If there's a clear way to handle cases like these, we'll implement it!

The answer may ultimately be that we can only represent Hessians for array-input array-output functions, in which case flattening a container-input function into an array-input one, then working with the Hessian of the flattened function, may be the right answer. JAX could provide utilities for doing that (we had [nice ones in Autograd, with a slick implementation](https://github.com/HIPS/autograd/blob/master/autograd/misc/flatten.py)).

In the meantime, flattening things yourself probably makes the most sense, unless you want to wait a few days for JAX to gain some flattening utilities. An alternative might be to use a [Lanczos iteration](https://en.wikipedia.org/wiki/Lanczos_algorithm) together with a Hessian-vector product, which you can express easily in JAX. Then you'd only have to deal with vectors, rather than having to worry about how to represent matrices, and we know how to handle tuples/lists/dicts there. (But Lanczos would only be accurate for extremal eigenvalues, and its numerical effectiveness would depend on the conditioning of the Hessian, whereas direct eigenvalue algorithms would be independent of the conditioning.)
> I really mean that as a question! Do you have a clear sense of what would make sense from your perspective? If there's a clear way to handle cases like these, we'll implement it!

I'm exploring JAX as a more straightforward way to do higher order derivatives. Right now I am using tensorflow (`tf.hessians`) but it quickly becomes clunky and it doesn't work in eager mode. My real use case is to do some analysis on the eigenvalues and eigenvectors of a neural network with a few thousand parameters.

> The answer may ultimately be that we can only represent Hessians for array-input array-output functions, in which case flattening a container-input function into an array-input one, then working with the Hessian of the flattened function, may be the right answer. JAX could provide utilities for doing that (we had nice ones in Autograd, with a slick implementation).

Flattening is currently the most straightforward way to do this in pytorch (which also has nice flattening utilities) say when constructing fisher information matrices in RL. Unfortunately, flattening in tensorflow graph mode cannot be used with `tf.hessians` unless the flattened version was used to predict the output of the network.

I think the most useful way for hessians of arbitrary matrices or compositions on matrices would be:
1. Have weight matrices represented as `[Input_i x Output_i]` for `i` the index of the layer.
2. Get the output of loss function by composing the weight matrices and required non-linearities.
3. Flatten the matrices into one long array and (somehow) use `jax.hessians` wrto the computed loss to get a matrix of size `[sum(prod(Input_i, Output_i)) x sum(prod(Input_i, Output_i))]`.



> An alternative might be to use a Lanczos iteration together with a Hessian-vector product, which you can express easily in JAX. Then you'd only have to deal with vectors, rather than having to worry about how to represent matrices, and we know how to handle tuples/lists/dicts there. (But Lanczos would only be accurate for extremal eigenvalues, and its numerical effectiveness would depend on the conditioning of the Hessian, whereas direct eigenvalue algorithms would be independent of the conditioning.)

I really want to compute the eigenspectrum (density of eigenvalues) and associated eigendirections so accuracy is important.

We had some conversations about this, and I think our plan is to:
1. add flattening utilities (track it in #190)
2. generalize `hessian` to work on containers (we chose a representation to go with).
I'm glad #201 didn't close this issue, because it seems to have broken `hessian`! I'm looking at it now.

Resolving patch:
diff --git a/examples/advi.py b/examples/advi.py
--- a/examples/advi.py
+++ b/examples/advi.py
@@ -1,3 +1,17 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """Automatic differentiation variational inference in Numpy and JAX.
 
 This demo fits a Gaussian approximation to an intractable, unnormalized
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -180,6 +180,7 @@ def jacfun(*args, **kwargs):
     return tree_transpose(tree_structure(example_args), tree_structure(y), jac)
 
   return jacfun
+jacobian = jacrev
 
 def hessian(fun):
   return jacfwd(jacrev(fun))
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -27,7 +27,7 @@
 from .. import core
 from ..core import Trace, Tracer, new_master, pack, AbstractTuple, JaxTuple
 from ..abstract_arrays import ShapedArray, make_shaped_array, array_types
-from ..ad_util import add_jaxvals_p
+from ..ad_util import add_jaxvals_p, zeros_like_p, zeros_like_jaxval
 from ..linear_util import transformation, transformation_with_aux, wrap_init
 from ..tree_util import register_pytree_node
 from ..util import unzip2, partial, safe_map
@@ -228,6 +228,8 @@ def reducer_batcher(prim, batched_args, batch_dims, axes, **kwargs):
     kwargs['input_shape'] = operand.shape
   return prim.bind(operand, axes=axes, **kwargs), bdim_out
 
+# set up primitive batches for ad_util primitives
+
 def add_batched(batched_args, batch_dims):
   bdx, bdy = batch_dims
   if bdx == bdy:
@@ -238,6 +240,12 @@ def add_batched(batched_args, batch_dims):
     return add_jaxvals_p.bind(xs, ys), 0
 primitive_batchers[add_jaxvals_p] = add_batched
 
+def zeros_like_batched(batched_args, batch_dims):
+  val, = batched_args
+  bdim, = batch_dims
+  return zeros_like_jaxval(val), bdim
+primitive_batchers[zeros_like_p] = zeros_like_batched
+
 
 ### util
 
@@ -287,6 +295,9 @@ def moveaxis(sz, dst, src, x):
       return pack(map(partial(moveaxis, sz), dst, src, x))
     elif type(src) is tuple:
       return pack(map(partial(moveaxis, sz, dst), src, x))
+    elif type(dst) is tuple:
+      srcs = (src,) * len(dst)
+      return pack(map(partial(moveaxis, sz), dst, srcs, x))
     else:
       return pack(map(partial(moveaxis, sz, dst, src), x))
   elif isinstance(aval, ShapedArray):


Test patch:
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -427,6 +427,11 @@ def assertAllClose(self, x, y, check_dtypes, atol=None, rtol=None):
       self.assertEqual(len(x), len(y))
       for x_elt, y_elt in zip(x, y):
         self.assertAllClose(x_elt, y_elt, check_dtypes, atol=atol, rtol=rtol)
+    elif isinstance(x, dict):
+      self.assertIsInstance(y, dict)
+      self.assertEqual(set(x.keys()), set(y.keys()))
+      for k in x.keys():
+        self.assertAllClose(x[k], y[k], check_dtypes, atol=atol, rtol=rtol)
     else:
       is_array = lambda x: hasattr(x, '__array__') or onp.isscalar(x)
       self.assertTrue(is_array(x))
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -23,7 +23,7 @@
 from jax import test_util as jtu
 
 import jax.numpy as np
-from jax import jit, grad, device_get, device_put, jacfwd, jacrev
+from jax import jit, grad, device_get, device_put, jacfwd, jacrev, hessian
 from jax import api
 from jax.core import Primitive
 from jax.interpreters.partial_eval import def_abstract_eval
@@ -260,6 +260,15 @@ def test_jacobian(self):
     f = lambda x: np.tanh(np.dot(A, x))
     assert onp.allclose(jacfwd(f)(x), jacrev(f)(x))
 
+  @jtu.skip_on_devices("tpu")
+  def test_hessian(self):
+    R = onp.random.RandomState(0).randn
+    A = R(4, 4)
+    x = R(4)
+
+    f = lambda x: np.dot(x, np.dot(A, x))
+    assert onp.allclose(hessian(f)(x), A + A.T)
+
   def test_std_basis(self):
     basis = api._std_basis(np.zeros(3))
     assert getattr(basis, "shape", None) == (3, 3)
@@ -276,6 +285,42 @@ def test_std_basis(self):
     assert getattr(basis[1][0], "shape", None) == (16, 3)
     assert getattr(basis[1][1], "shape", None) == (16, 3, 4)
 
+  @jtu.skip_on_devices("tpu")
+  def test_jacobian_on_pytrees(self):
+    for jacfun in [jacfwd, jacrev]:
+      ans = jacfun(lambda x, y: (x, y))(0., 1.)
+      expected = (1., 0.)
+      self.assertAllClose(ans, expected, check_dtypes=False)
+
+      ans = jacfun(lambda x, y: (x, y), 1)(0., 1.)
+      expected = (0., 1.)
+      self.assertAllClose(ans, expected, check_dtypes=False)
+
+      ans = jacfun(lambda x, y: (x, y), (0, 1))(0., 1.)
+      expected = ((1., 0.),
+                  (0., 1.),)
+      self.assertAllClose(ans, expected, check_dtypes=False)
+
+      ans = jacfun(lambda x: x[:2])((1., 2., 3.))
+      expected = ((1., 0., 0.),
+                  (0., 1., 0.))
+      self.assertAllClose(ans, expected, check_dtypes=False)
+
+      R = onp.random.RandomState(0).randn
+      x = R(2)
+      y = R(3)
+      ans = jacfun(lambda x, y: {'x': x, 'xy': np.outer(x, y)})(x, y)
+      expected = {'x': onp.eye(2),
+                  'xy': onp.kron(onp.eye(2), y[:, None]).reshape(2, 3, 2)}
+      self.assertAllClose(ans, expected, check_dtypes=False)
+
+  @jtu.skip_on_devices("tpu")
+  def test_hessian_on_pytrees(self):
+    ans = hessian(lambda x: np.array(x)**2)((1., 2.))
+    expected = ((onp.array([2., 0.]), onp.array([0., 0.])),
+                (onp.array([0., 0.]), onp.array([0., 2.])))
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
 
 if __name__ == '__main__':
   absltest.main()