Repository: pyro-ppl/numpyro

Issue:
[Feature request] hessian(f)(x) = jacrev(jacrev(f))(x))
Hi! fehiepsi suggested an edit to the Laplace Approximation to get an uncertainty estimate when using Minimize() results for a system of ODEs (solved with jax.odeint, which has a custom_vjp). 

Specifically I'm currently getting the following error: FilteredStackTrace: TypeError: can't apply forward-mode autodiff (jvp) to a custom_vjp function.

The suggested edit:  hessian(f)(x) = jacfwd(jacrev(f))(x))   -->  hessian(f)(x) = jacrev(jacrev(f))(x)
(https://forum.pyro.ai/t/making-a-for-loop-more-efficient/3437/18)


Discussion:

Resolving patch:
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
--- a/numpyro/infer/autoguide.py
+++ b/numpyro/infer/autoguide.py
@@ -1086,8 +1086,33 @@ class AutoLaplaceApproximation(AutoContinuous):
 
         guide = AutoLaplaceApproximation(model, ...)
         svi = SVI(model, guide, ...)
+
+    :param callable hessian_fn: EXPERIMENTAL a function that takes a function `f`
+        and a vector `x`and returns the hessian of `f` at `x`. By default, we use
+        ``lambda f, x: jax.hessian(f)(x)``. Other alternatives can be
+        ``lambda f, x: jax.jacobian(jax.jacobian(f))(x)`` or
+        ``lambda f, x: jax.hessian(f)(x) + 1e-3 * jnp.eye(x.shape[0])``. The later
+        example is helpful when the hessian of `f` at `x` is not positive definite.
+        Note that the output hessian is the precision matrix of the laplace
+        approximation.
     """
 
+    def __init__(
+        self,
+        model,
+        *,
+        prefix="auto",
+        init_loc_fn=init_to_uniform,
+        create_plates=None,
+        hessian_fn=None,
+    ):
+        super().__init__(
+            model, prefix=prefix, init_loc_fn=init_loc_fn, create_plates=create_plates
+        )
+        self._hessian_fn = (
+            hessian_fn if hessian_fn is not None else (lambda f, x: hessian(f)(x))
+        )
+
     def _setup_prototype(self, *args, **kwargs):
         super(AutoLaplaceApproximation, self)._setup_prototype(*args, **kwargs)
 
@@ -1114,7 +1139,7 @@ def loss_fn(z):
             return self._loss_fn(params1)
 
         loc = params["{}_loc".format(self.prefix)]
-        precision = hessian(loss_fn)(loc)
+        precision = self._hessian_fn(loss_fn, loc)
         scale_tril = cholesky_of_inverse(precision)
         if not_jax_tracer(scale_tril):
             if np.any(np.isnan(scale_tril)):


Test patch:
diff --git a/test/infer/test_autoguide.py b/test/infer/test_autoguide.py
--- a/test/infer/test_autoguide.py
+++ b/test/infer/test_autoguide.py
@@ -6,7 +6,7 @@
 from numpy.testing import assert_allclose
 import pytest
 
-from jax import jit, lax, random
+from jax import jacobian, jit, lax, random
 from jax.experimental.stax import Dense
 import jax.numpy as jnp
 from jax.test_util import check_eq
@@ -361,6 +361,23 @@ def model(x, y):
         guide.sample_posterior(random.PRNGKey(1), params)
 
 
+def test_laplace_approximation_custom_hessian():
+    def model(x, y):
+        a = numpyro.sample("a", dist.Normal(0, 10))
+        b = numpyro.sample("b", dist.Normal(0, 10))
+        mu = a + b * x
+        numpyro.sample("y", dist.Normal(mu, 1), obs=y)
+
+    x = random.normal(random.PRNGKey(0), (100,))
+    y = 1 + 2 * x
+    guide = AutoLaplaceApproximation(
+        model, hessian_fn=lambda f, x: jacobian(jacobian(f))(x)
+    )
+    svi = SVI(model, guide, optim.Adam(0.1), Trace_ELBO(), x=x, y=y)
+    svi_result = svi.run(random.PRNGKey(0), 10000, progress_bar=False)
+    guide.get_transform(svi_result.params)
+
+
 def test_improper():
     y = random.normal(random.PRNGKey(0), (100,))