Add support for Optax optimizers (#1034)

* Add support for Optax optimizers

* Test Optax

* Fix isort linter

* Add optax to docs requirements

* Simplify optax wrapper

* Better handling of import error

* Pin Sphinx version

* Address comments

* Add optax_to_numpyro to docs

* isort

* Edit optax snippet in SVI docstring

* Update docs

* Use optax.OptState for type hints

* Allow jax.experimental.optimizers.Optimizer instances to be used directly in SVI

* Rerun CI

* Pin Jinja2 version

Author: tcbegley

docs/requirements.txt

@@ -3,7 +3,8 @@ flax
 funsor
 jax>=0.1.65
 jaxlib>=0.1.45
-nbsphinx==0.8.1
+optax==0.0.6
+nbsphinx>=0.8.4
 sphinx-gallery
 tfp-nightly  # TODO: change this to tensorflow-probability when it is stable
 tqdm

docs/source/optimizers.rst

@@ -64,4 +64,8 @@ SM3
 .. autoclass:: numpyro.optim.SM3
    :members:
    :undoc-members:
-   :inherited-members:
+   :inherited-members:
+
+Optax support
+-------------
+.. autofunction:: numpyro.contrib.optim.optax_to_numpyro

numpyro/contrib/optim.py

@@ -0,0 +1,46 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+This module provides a wrapper for Optax optimizers so that they can be used with
+NumPyro inference algorithms.
+"""
+
+from typing import Tuple, TypeVar
+
+import optax
+
+from numpyro.optim import _NumPyroOptim
+
+_Params = TypeVar("_Params")
+_State = Tuple[_Params, optax.OptState]
+
+
+def optax_to_numpyro(transformation: optax.GradientTransformation) -> _NumPyroOptim:
+    """
+    This function produces a ``numpyro.optim._NumPyroOptim`` instance from an
+    ``optax.GradientTransformation`` so that it can be used with
+    ``numpyro.infer.svi.SVI``. It is a lightweight wrapper that recreates the
+    ``(init_fn, update_fn, get_params_fn)`` interface defined by
+    :mod:`jax.experimental.optimizers`.
+
+    :param transformation: An ``optax.GradientTransformation`` instance to wrap.
+    :return: An instance of ``numpyro.optim._NumPyroOptim`` wrapping the supplied
+        Optax optimizer.
+    """
+
+    def init_fn(params: _Params) -> _State:
+        opt_state = transformation.init(params)
+        return params, opt_state
+
+    def update_fn(step, grads: _Params, state: _State) -> _State:
+        params, opt_state = state
+        updates, opt_state = transformation.update(grads, opt_state, params)
+        updated_params = optax.apply_updates(params, updates)
+        return updated_params, opt_state
+
+    def get_params_fn(state: _State) -> _Params:
+        params, _ = state
+        return params
+
+    return _NumPyroOptim(lambda x, y, z: (x, y, z), init_fn, update_fn, get_params_fn)

numpyro/infer/svi.py

@@ -6,6 +6,7 @@
 
 import tqdm
 
+import jax
 from jax import jit, lax, random
 import jax.numpy as jnp
 from jax.tree_util import tree_map
@@ -14,6 +15,7 @@
 from numpyro.distributions.transforms import biject_to
 from numpyro.handlers import replay, seed, trace
 from numpyro.infer.util import transform_fn
+from numpyro.optim import _NumPyroOptim
 
 SVIState = namedtuple("SVIState", ["optim_state", "rng_key"])
 """
@@ -80,7 +82,14 @@ class SVI(object):
     :param model: Python callable with Pyro primitives for the model.
     :param guide: Python callable with Pyro primitives for the guide
         (recognition network).
-    :param optim: an instance of :class:`~numpyro.optim._NumpyroOptim`.
+    :param optim: An instance of :class:`~numpyro.optim._NumpyroOptim`, a
+        ``jax.experimental.optimizers.Optimizer`` or an Optax
+        ``GradientTransformation``. If you pass an Optax optimizer it will
+        automatically be wrapped using :func:`numpyro.contrib.optim.optax_to_numpyro`.
+
+            >>> from optax import adam, chain, clip
+            >>> svi = SVI(model, guide, chain(clip(10.0), adam(1e-3)), loss=Trace_ELBO())
+
     :param loss: ELBO loss, i.e. negative Evidence Lower Bound, to minimize.
     :param static_kwargs: static arguments for the model / guide, i.e. arguments
         that remain constant during fitting.
@@ -91,10 +100,36 @@ def __init__(self, model, guide, optim, loss, **static_kwargs):
         self.model = model
         self.guide = guide
         self.loss = loss
-        self.optim = optim
         self.static_kwargs = static_kwargs
         self.constrain_fn = None
 
+        if isinstance(optim, _NumPyroOptim):
+            self.optim = optim
+        elif isinstance(optim, jax.experimental.optimizers.Optimizer):
+            self.optim = _NumPyroOptim(lambda *args: args, *optim)
+        else:
+            try:
+                import optax
+
+                from numpyro.contrib.optim import optax_to_numpyro
+            except ImportError:
+                raise ImportError(
+                    "It looks like you tried to use an optimizer that isn't an "
+                    "instance of numpyro.optim._NumPyroOptim or "
+                    "jax.experimental.optimizers.Optimizer. There is experimental "
+                    "support for Optax optimizers, but you need to install Optax. "
+                    "It can be installed with `pip install optax`."
+                )
+
+            if not isinstance(optim, optax.GradientTransformation):
+                raise TypeError(
+                    "Expected either an instance of numpyro.optim._NumPyroOptim, "
+                    "jax.experimental.optimizers.Optimizer or "
+                    "optax.GradientTransformation. Got {}".format(type(optim))
+                )
+
+            self.optim = optax_to_numpyro(optim)
+
     def init(self, rng_key, *args, **kwargs):
         """
         Gets the initial SVI state.

setup.cfg

@@ -8,7 +8,7 @@ profile = black
 skip_glob = .ipynb_checkpoints
 known_first_party = funsor, numpyro, test
 known_third_party = opt_einsum
-known_jax = flax, haiku, jax, tensorflow_probability
+known_jax = flax, haiku, jax, optax, tensorflow_probability
 sections = FUTURE, STDLIB, THIRDPARTY, JAX, FIRSTPARTY, LOCALFOLDER
 force_sort_within_sections = true
 combine_as_imports = true

setup.py

@@ -39,8 +39,9 @@
     extras_require={
         "doc": [
             "ipython",  # sphinx needs this to render codes
+            "jinja2<3.0.0",
             "nbsphinx",
-            "sphinx",
+            "sphinx<4.0.0",
             "sphinx_rtd_theme",
             "sphinx-gallery",
         ],
@@ -58,6 +59,7 @@
             # TODO: bump funsor version before the release
             "funsor @ git+https://github.com/pyro-ppl/funsor.git@d5574988665dd822ec64e41f2b54b9dc929959dc",
             "graphviz",
+            "optax==0.0.6",
             # TODO: change this to tensorflow_probability>0.12.1 when the next version
             # of tfp is released. The current release is not compatible with jax>=0.2.12.
             "tfp-nightly",

test/contrib/test_optim.py

@@ -0,0 +1,114 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+from numpy.testing import assert_allclose
+import pytest
+
+from jax import grad, jit, partial, random
+from jax.lax import fori_loop
+import jax.numpy as jnp
+from jax.test_util import check_close
+
+import numpyro
+import numpyro.distributions as dist
+from numpyro.distributions import constraints
+from numpyro.infer import SVI, RenyiELBO, Trace_ELBO
+
+try:
+    import optax
+
+    from numpyro.contrib.optim import optax_to_numpyro
+
+    # the optimizer test is parameterized by different optax optimizers, but we have
+    # to define them here to ensure that `optax` is defined. pytest.mark.parameterize
+    # decorators are run even if tests are skipped at the top of the file.
+    optimizers = [
+        (optax.adam, (1e-2,), {}),
+        # clipped adam
+        (optax.chain, (optax.clip(10.0), optax.adam(1e-2)), {}),
+        (optax.adagrad, (1e-1,), {}),
+        # SGD with momentum
+        (optax.sgd, (1e-2,), {"momentum": 0.9}),
+        (optax.rmsprop, (1e-2,), {"decay": 0.95}),
+        # RMSProp with momentum
+        (optax.rmsprop, (1e-4,), {"decay": 0.9, "momentum": 0.9}),
+        (optax.sgd, (1e-2,), {}),
+    ]
+except ImportError:
+    pytestmark = pytest.mark.skip(reason="optax is not installed")
+    optimizers = []
+
+
+def loss(params):
+    return jnp.sum(params["x"] ** 2 + params["y"] ** 2)
+
+
+@partial(jit, static_argnums=(1,))
+def step(opt_state, optim):
+    params = optim.get_params(opt_state)
+    g = grad(loss)(params)
+    return optim.update(g, opt_state)
+
+
+@pytest.mark.parametrize("optim_class, args, kwargs", optimizers)
+def test_optim_multi_params(optim_class, args, kwargs):
+    params = {"x": jnp.array([1.0, 1.0, 1.0]), "y": jnp.array([-1, -1.0, -1.0])}
+    opt = optax_to_numpyro(optim_class(*args, **kwargs))
+    opt_state = opt.init(params)
+    for i in range(2000):
+        opt_state = step(opt_state, opt)
+    for _, param in opt.get_params(opt_state).items():
+        assert jnp.allclose(param, jnp.zeros(3))
+
+
+@pytest.mark.parametrize("elbo", [Trace_ELBO(), RenyiELBO(num_particles=10)])
+def test_beta_bernoulli(elbo):
+    data = jnp.array([1.0] * 8 + [0.0] * 2)
+
+    def model(data):
+        f = numpyro.sample("beta", dist.Beta(1.0, 1.0))
+        numpyro.sample("obs", dist.Bernoulli(f), obs=data)
+
+    def guide(data):
+        alpha_q = numpyro.param("alpha_q", 1.0, constraint=constraints.positive)
+        beta_q = numpyro.param("beta_q", 1.0, constraint=constraints.positive)
+        numpyro.sample("beta", dist.Beta(alpha_q, beta_q))
+
+    adam = optax.adam(0.05)
+    svi = SVI(model, guide, adam, elbo)
+    svi_state = svi.init(random.PRNGKey(1), data)
+    assert_allclose(svi.optim.get_params(svi_state.optim_state)["alpha_q"], 0.0)
+
+    def body_fn(i, val):
+        svi_state, _ = svi.update(val, data)
+        return svi_state
+
+    svi_state = fori_loop(0, 2000, body_fn, svi_state)
+    params = svi.get_params(svi_state)
+    assert_allclose(
+        params["alpha_q"] / (params["alpha_q"] + params["beta_q"]),
+        0.8,
+        atol=0.05,
+        rtol=0.05,
+    )
+
+
+def test_jitted_update_fn():
+    data = jnp.array([1.0] * 8 + [0.0] * 2)
+
+    def model(data):
+        f = numpyro.sample("beta", dist.Beta(1.0, 1.0))
+        numpyro.sample("obs", dist.Bernoulli(f), obs=data)
+
+    def guide(data):
+        alpha_q = numpyro.param("alpha_q", 1.0, constraint=constraints.positive)
+        beta_q = numpyro.param("beta_q", 1.0, constraint=constraints.positive)
+        numpyro.sample("beta", dist.Beta(alpha_q, beta_q))
+
+    adam = optax.adam(0.05)
+    svi = SVI(model, guide, adam, Trace_ELBO())
+    svi_state = svi.init(random.PRNGKey(1), data)
+    expected = svi.get_params(svi.update(svi_state, data)[0])
+
+    actual = svi.get_params(jit(svi.update)(svi_state, data=data)[0])
+    check_close(actual, expected, atol=1e-5)

test/infer/test_svi.py

@@ -4,6 +4,7 @@
 from numpy.testing import assert_allclose
 import pytest
 
+import jax
 from jax import jit, random, value_and_grad
 import jax.numpy as jnp
 from jax.test_util import check_close
@@ -41,7 +42,10 @@ def renyi_loss_fn(x):
 
 
 @pytest.mark.parametrize("elbo", [Trace_ELBO(), RenyiELBO(num_particles=10)])
-def test_beta_bernoulli(elbo):
+@pytest.mark.parametrize(
+    "optimizer", [optim.Adam(0.05), jax.experimental.optimizers.adam(0.05)]
+)
+def test_beta_bernoulli(elbo, optimizer):
     data = jnp.array([1.0] * 8 + [0.0] * 2)
 
     def model(data):
@@ -53,10 +57,9 @@ def guide(data):
         beta_q = numpyro.param("beta_q", 1.0, constraint=constraints.positive)
         numpyro.sample("beta", dist.Beta(alpha_q, beta_q))
 
-    adam = optim.Adam(0.05)
-    svi = SVI(model, guide, adam, elbo)
+    svi = SVI(model, guide, optimizer, elbo)
     svi_state = svi.init(random.PRNGKey(1), data)
-    assert_allclose(adam.get_params(svi_state.optim_state)["alpha_q"], 0.0)
+    assert_allclose(svi.optim.get_params(svi_state.optim_state)["alpha_q"], 0.0)
 
     def body_fn(i, val):
         svi_state, _ = svi.update(val, data)