Repository: pyro-ppl/numpyro

Issue:
Support jax.scipy minimizers in NumPyro
After jax 0.1.74, `bfgs` minimizer is available in JAX. It would be nice to have a BFGS wrapper in `numpyro.optim`.

In jax/scipy, 'BFGS' is provided through `method` keyword in [minimize](https://docs.scipy.org/doc/scipy/reference/optimize.minimize-bfgs.html#optimize-minimize-bfgs). I am not sure if we should make a wrapper `numpyro.optim.Minimize` for `minimize` or we should split those methods into separated classes such as `numpyro.optim.BFGS`. With the former, new minimizers in JAX will be supported in NumPyro without modifications. With the latter, it is more explicit and friendlier to PyTorch users. Either way looks good to me.

cc @Joshuaalbert who implemented BFGS minimizer in JAX


Discussion:
I like the first way. It'll be familiar for people who have used scipy, where you also specify a method name. And then as you noted, it'll let people use new methods as they become available. JAX is following the scipy interface so you can expect the output of minimize to stay consistent.
Yeah, I agree. After some tfp prs, if no one wants to tackle this issue, I'll do it.

Resolving patch:
diff --git a/numpyro/infer/svi.py b/numpyro/infer/svi.py
--- a/numpyro/infer/svi.py
+++ b/numpyro/infer/svi.py
@@ -3,7 +3,7 @@
 
 from functools import namedtuple, partial
 
-from jax import random, value_and_grad
+from jax import random
 
 from numpyro.distributions import constraints
 from numpyro.distributions.transforms import biject_to
@@ -18,6 +18,11 @@
 """
 
 
+def _apply_loss_fn(loss_fn, rng_key, constrain_fn, model, guide,
+                   args, kwargs, static_kwargs, params):
+    return loss_fn(rng_key, constrain_fn(params), model, guide, *args, **kwargs, **static_kwargs)
+
+
 class SVI(object):
     """
     Stochastic Variational Inference given an ELBO loss objective.
@@ -127,11 +132,9 @@ def update(self, svi_state, *args, **kwargs):
         :return: tuple of `(svi_state, loss)`.
         """
         rng_key, rng_key_step = random.split(svi_state.rng_key)
-        params = self.optim.get_params(svi_state.optim_state)
-        loss_val, grads = value_and_grad(
-            lambda x: self.loss.loss(rng_key_step, self.constrain_fn(x), self.model, self.guide,
-                                     *args, **kwargs, **self.static_kwargs))(params)
-        optim_state = self.optim.update(grads, svi_state.optim_state)
+        loss_fn = partial(_apply_loss_fn, self.loss.loss, rng_key_step, self.constrain_fn, self.model,
+                          self.guide, args, kwargs, self.static_kwargs)
+        loss_val, optim_state = self.optim.eval_and_update(loss_fn, svi_state.optim_state)
         return SVIState(optim_state, rng_key), loss_val
 
     def evaluate(self, svi_state, *args, **kwargs):
diff --git a/numpyro/optim.py b/numpyro/optim.py
--- a/numpyro/optim.py
+++ b/numpyro/optim.py
@@ -7,16 +7,21 @@
 suited for working with NumPyro inference algorithms.
 """
 
+from collections import namedtuple
 from typing import Callable, Tuple, TypeVar
 
+from jax import value_and_grad
 from jax.experimental import optimizers
+from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
-from jax.tree_util import tree_map
+from jax.scipy.optimize import minimize
+from jax.tree_util import register_pytree_node, tree_map
 
 __all__ = [
     'Adam',
     'Adagrad',
     'ClippedAdam',
+    'Minimize',
     'Momentum',
     'RMSProp',
     'RMSPropMomentum',
@@ -29,7 +34,7 @@
 _IterOptState = Tuple[int, _OptState]
 
 
-class _NumpyroOptim(object):
+class _NumPyroOptim(object):
     def __init__(self, optim_fn: Callable, *args, **kwargs) -> None:
         self.init_fn, self.update_fn, self.get_params_fn = optim_fn(*args, **kwargs)
 
@@ -55,6 +60,23 @@ def update(self, g: _Params, state: _IterOptState) -> _IterOptState:
         opt_state = self.update_fn(i, g, opt_state)
         return i + 1, opt_state
 
+    def eval_and_update(self, fn: Callable, state: _IterOptState) -> _IterOptState:
+        """
+        Performs an optimization step for the objective function `fn`.
+        For most optimizers, the update is performed based on the gradient
+        of the objective function w.r.t. the current state. However, for
+        some optimizers such as :class:`Minimize`, the update is performed
+        by reevaluating the function multiple times to get optimal
+        parameters.
+
+        :param fn: objective function.
+        :param state: current optimizer state.
+        :return: a pair of the output of objective function and the new optimizer state.
+        """
+        params = self.get_params(state)
+        out, grads = value_and_grad(fn)(params)
+        return out, self.update(grads, state)
+
     def get_params(self, state: _IterOptState) -> _Params:
         """
         Get current parameter values.
@@ -76,12 +98,12 @@ def _wrapped(cls):
 
 
 @_add_doc(optimizers.adam)
-class Adam(_NumpyroOptim):
+class Adam(_NumPyroOptim):
     def __init__(self, *args, **kwargs):
         super(Adam, self).__init__(optimizers.adam, *args, **kwargs)
 
 
-class ClippedAdam(_NumpyroOptim):
+class ClippedAdam(_NumPyroOptim):
     """
     :class:`~numpyro.optim.Adam` optimizer with gradient clipping.
 
@@ -106,36 +128,117 @@ def update(self, g, state):
 
 
 @_add_doc(optimizers.adagrad)
-class Adagrad(_NumpyroOptim):
+class Adagrad(_NumPyroOptim):
     def __init__(self, *args, **kwargs):
         super(Adagrad, self).__init__(optimizers.adagrad, *args, **kwargs)
 
 
 @_add_doc(optimizers.momentum)
-class Momentum(_NumpyroOptim):
+class Momentum(_NumPyroOptim):
     def __init__(self, *args, **kwargs):
         super(Momentum, self).__init__(optimizers.momentum, *args, **kwargs)
 
 
 @_add_doc(optimizers.rmsprop)
-class RMSProp(_NumpyroOptim):
+class RMSProp(_NumPyroOptim):
     def __init__(self, *args, **kwargs):
         super(RMSProp, self).__init__(optimizers.rmsprop, *args, **kwargs)
 
 
 @_add_doc(optimizers.rmsprop_momentum)
-class RMSPropMomentum(_NumpyroOptim):
+class RMSPropMomentum(_NumPyroOptim):
     def __init__(self, *args, **kwargs):
         super(RMSPropMomentum, self).__init__(optimizers.rmsprop_momentum, *args, **kwargs)
 
 
 @_add_doc(optimizers.sgd)
-class SGD(_NumpyroOptim):
+class SGD(_NumPyroOptim):
     def __init__(self, *args, **kwargs):
         super(SGD, self).__init__(optimizers.sgd, *args, **kwargs)
 
 
 @_add_doc(optimizers.sm3)
-class SM3(_NumpyroOptim):
+class SM3(_NumPyroOptim):
     def __init__(self, *args, **kwargs):
         super(SM3, self).__init__(optimizers.sm3, *args, **kwargs)
+
+
+# TODO: currently, jax.scipy.optimize.minimize only supports 1D input,
+# so we need to add the following mechanism to transform params to flat_params
+# and pass `unravel_fn` arround.
+# When arbitrary pytree is supported in JAX, we can just simply use
+# identity functions for `init_fn` and `get_params`.
+_MinimizeState = namedtuple("MinimizeState", ["flat_params", "unravel_fn"])
+register_pytree_node(
+    _MinimizeState,
+    lambda state: ((state.flat_params,), (state.unravel_fn,)),
+    lambda data, xs: _MinimizeState(xs[0], data[0]))
+
+
+def _minimize_wrapper():
+    def init_fn(params):
+        flat_params, unravel_fn = ravel_pytree(params)
+        return _MinimizeState(flat_params, unravel_fn)
+
+    def update_fn(i, grad_tree, opt_state):
+        # we don't use update_fn in Minimize, so let it do nothing
+        return opt_state
+
+    def get_params(opt_state):
+        flat_params, unravel_fn = opt_state
+        return unravel_fn(flat_params)
+
+    return init_fn, update_fn, get_params
+
+
+class Minimize(_NumPyroOptim):
+    """
+    Wrapper class for the JAX minimizer: :func:`~jax.scipy.optimize.minimize`.
+
+    .. warnings: This optimizer is intended to be used with static guides such
+        as empty guides (maximum likelihood estimate), delta guides (MAP estimate),
+        or :class:`~numpyro.infer.autoguide.AutoLaplaceApproximation`.
+        Using this in stochastic setting is either expensive or hard to converge.
+
+    **Example:**
+
+    .. doctest::
+
+        >>> from numpy.testing import assert_allclose
+        >>> from jax import random
+        >>> import jax.numpy as jnp
+        >>> import numpyro
+        >>> import numpyro.distributions as dist
+        >>> from numpyro.infer import SVI, ELBO
+        >>> from numpyro.infer.autoguide import AutoLaplaceApproximation
+
+        >>> def model(x, y):
+        ...     a = numpyro.sample("a", dist.Normal(0, 1))
+        ...     b = numpyro.sample("b", dist.Normal(0, 1))
+        ...     with numpyro.plate("N", y.shape[0]):
+        ...         numpyro.sample("obs", dist.Normal(a + b * x, 0.1), obs=y)
+
+        >>> x = jnp.linspace(0, 10, 100)
+        >>> y = 3 * x + 2
+        >>> optimizer = numpyro.optim.Minimize()
+        >>> guide = AutoLaplaceApproximation(model)
+        >>> svi = SVI(model, guide, optimizer, loss=ELBO())
+        >>> init_state = svi.init(random.PRNGKey(0), x, y)
+        >>> optimal_state, loss = svi.update(init_state, x, y)
+        >>> params = svi.get_params(optimal_state)  # get guide's parameters
+        >>> quantiles = guide.quantiles(params, 0.5)  # get means of posterior samples
+        >>> assert_allclose(quantiles["a"], 2., atol=1e-3)
+        >>> assert_allclose(quantiles["b"], 3., atol=1e-3)
+    """
+    def __init__(self, method="BFGS", **kwargs):
+        super().__init__(_minimize_wrapper)
+        self._method = method
+        self._kwargs = kwargs
+
+    def eval_and_update(self, fn: Callable, state: _IterOptState) -> _IterOptState:
+        i, (flat_params, unravel_fn) = state
+        results = minimize(lambda x: fn(unravel_fn(x)), flat_params, (),
+                           method=self._method, **self._kwargs)
+        flat_params, out = results.x, results.fun
+        state = (i + 1, _MinimizeState(flat_params, unravel_fn))
+        return out, state


Test patch: