Repository: pyro-ppl/numpyro

Issue:
Integrate with flax/haiku
Currently, we support [stax](https://jax.readthedocs.io/en/latest/jax.experimental.stax.html) but it lacks some popular layers such as RNN. Some other candidates include [flax](https://github.com/google/flax), [haiku](https://github.com/deepmind/dm-haiku). It might not take much effort to use those libraries in NumPyro.

We might start with some utilities/handlers in `numpyro.contrib.flax` or `numpyro.contrib.haiku`, add some examples (which illustrates how to set priors for some parameters of neural nets), and gradually enhance the API.

Some references:
+ [lift](http://num.pyro.ai/en/stable/handlers.html#lift) handler: converts a `param` site to a `sample` site
+ [module](http://num.pyro.ai/en/stable/primitives.html#module) primitive: move nn's parameters to `param` statement
+ [VAE example](http://pyro.ai/numpyro/examples/vae.html)


Discussion:
Can i work on this ? Promise i'll finish the PR timely :D 
Hi @TuanNguyen27! It is great to hear that you are interested in this PR. Please take your time.

FYI, if you come up with several design choices and don't know which one is good, you can create a design doc (you can take a look a some [design docs](https://github.com/pyro-ppl/pyro/wiki/Design-Docs) that we created previously) and invite us to review/comment.
@fehiepsi For my first attempt, I will try to migrate VAE example to use flax/haiku. If that works, how do I further integrate this library into numpyro, do i re-implement some existing handlers to work with haiku, for example ? 
@TuanNguyen27 I think you can put required utilities/handlers in `contrib.flax`. Users will need to install flax/haiku to use the utilities/handlers in `contrib.flax`. You can add an ImportError for it, for example like [this](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/contrib/funsor/__init__.py#L5).
@fehiepsi I'm trying to rewrite `numpyro.module` to work with haiku's `Module` class, but somehow my losses come out as Nan during training. I'm wondering if you could take a look to see if there's anything i missed. 

https://gist.github.com/TuanNguyen27/57ab8264632ccded683ff819acb99253

[Runnable colab link here](https://colab.research.google.com/gist/TuanNguyen27/57ab8264632ccded683ff819acb99253/haiku_prototype.ipynb)

I had to supply rng keys every time `.apply` is called, and stax's `apply` doesn't require this, so I'm not sure how things translate between the two libraries.
I'm running into the same problem when using flax's `nn.Module`. I'd love to hear any suggestion / debugging tips to make more progress :) 

https://colab.research.google.com/gist/TuanNguyen27/409192b97d01e1ea5583f9e4499d306a/flax_prototype.ipynb
haiku's Module & flax's Module both require user to first initialize the parameter with rng & data input (data input could be completely unrelated to the real training data as long as it has the correct shape). When I reimplement `numpyro.module`, I initialized Encoder and Decoder's parameters with `(rng_key, jnp.ones(input_shape))`. I then wrap the initalized params with `numpyro.param(module_key, nn_params)` -- this is where I am not sure if the code works the same as `stax`. 

@TuanNguyen27 Awesome! You are so fast... Your code looks very reasonable to me. I will try to debug to see if I can isolate the nan issue for you.
@TuanNguyen27 It seems that you need to add some additional layers to the output of those nn modules:
```python
class Encoder(nn.Module):
    def apply(self, x, hidden_size = 512, latents = 10):
        ...
        return mean_x, jnp.exp(logvar_x)

class Decoder(nn.Module):
    def apply(self, z, hidden_size=512, output_shape=784):
        ...
        return nn.sigmoid(z)
```

While playing with your example, I really like the way you simplify the problem: we only need to make a new handler `module` and its definition is quite similar to the current `numpyro.module`. With that, to integrate with NumPyro, I think you only need to put it in a contrib file like `numpyro/contrib/nn.py` (probably make the name explicitly like `module` -> `stax_module`).

The reason for putting it into `contrib` is I am seeing there are further cool extensions from your work! It is important that those Flax parameters have corresponding names. This enables us to implement something similar to [random_module](http://docs.pyro.ai/en/stable/primitives.html#pyro.primitives.random_module) in Pyro, which turns an NN into a Bayesian NN. I remember some users in Pyro forum have asked for this feature. :)

cc @neerajprad I think you will be interested in @TuanNguyen27's work :)
that's great! I am also interested in converting a traditional NN into a Bayesian NN. In the past I have used [pyvarinf](https://github.com/ctallec/pyvarinf), which might be similar to how Pyro does the conversion. Let me rerun my notebook with your suggestion and make a PR :)

PS: Do you want both haiku and flax to integrate with numpyro or just one ? I can also try to see if Oryx has a similar module system.   

Resolving patch:
diff --git a/numpyro/contrib/module.py b/numpyro/contrib/module.py
new file mode 100644
--- /dev/null
+++ b/numpyro/contrib/module.py
@@ -0,0 +1,76 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+from functools import partial
+
+from jax import numpy as jnp
+
+import numpyro
+from numpyro.distributions.discrete import PRNGIdentity
+
+
+def flax_module(name, nn, input_shape=None):
+    """
+    Declare a :mod:`~flax` style neural network inside a
+    model so that its parameters are registered for optimization via
+    :func:`~numpyro.primitives.param` statements.
+
+    :param str name: name of the module to be registered.
+    :param flax.nn.Module nn: a `flax` Module which has .init and .apply methods
+    :param tuple input_shape: shape of the input taken by the
+        neural network.
+    :return: a callable with bound parameters that takes an array
+        as an input and returns the neural network transformed output
+        array.
+    """
+    try:
+        import flax  # noqa: F401
+    except ImportError:
+        raise ImportError("Looking like you want to use flax to declare "
+                          "nn modules. This is an experimental feature. "
+                          "You need to install `flax` to be able to use this feature. "
+                          "It can be installed with `pip install flax`.")
+    module_key = name + '$params'
+    nn_params = numpyro.param(module_key)
+    if nn_params is None:
+        if input_shape is None:
+            raise ValueError('Valid value for `input_shape` needed to initialize.')
+        # feed in dummy data to init params
+        rng_key = numpyro.sample(name + '$rng_key', PRNGIdentity())
+        _, nn_params = nn.init(rng_key, jnp.ones(input_shape))
+        numpyro.param(module_key, nn_params)
+    return partial(nn.call, nn_params)
+
+
+def haiku_module(name, nn, input_shape=None):
+    """
+    Declare a :mod:`~haiku` style neural network inside a
+    model so that its parameters are registered for optimization via
+    :func:`~numpyro.primitives.param` statements.
+
+    :param str name: name of the module to be registered.
+    :param haiku.Module nn: a `haiku` Module which has .init and .apply methods
+    :param tuple input_shape: shape of the input taken by the
+        neural network.
+    :return: a callable with bound parameters that takes an array
+        as an input and returns the neural network transformed output
+        array.
+    """
+    try:
+        import haiku  # noqa: F401
+    except ImportError:
+        raise ImportError("Looking like you want to use haiku to declare "
+                          "nn modules. This is an experimental feature. "
+                          "You need to install `haiku` to be able to use this feature. "
+                          "It can be installed with `pip install git+https://github.com/deepmind/dm-haiku`.")
+
+    module_key = name + '$params'
+    nn_params = numpyro.param(module_key)
+    if nn_params is None:
+        if input_shape is None:
+            raise ValueError('Valid value for `input_shape` needed to initialize.')
+        # feed in dummy data to init params
+        rng_key = numpyro.sample(name + '$rng_key', PRNGIdentity())
+        nn_params = nn.init(rng_key, jnp.ones(input_shape))
+        numpyro.param(module_key, nn_params)
+    return partial(nn.apply, nn_params, None)
diff --git a/numpyro/primitives.py b/numpyro/primitives.py
--- a/numpyro/primitives.py
+++ b/numpyro/primitives.py
@@ -195,7 +195,7 @@ def module(name, nn, input_shape=None):
     nn_params = param(module_key)
     if nn_params is None:
         if input_shape is None:
-            raise ValueError('Valid value for `input_size` needed to initialize.')
+            raise ValueError('Valid value for `input_shape` needed to initialize.')
         rng_key = numpyro.sample(name + '$rng_key', PRNGIdentity())
         _, nn_params = nn_init(rng_key, input_shape)
         param(module_key, nn_params)
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -50,6 +50,8 @@
             'funsor',
             'ipython',
             'isort',
+            'flax',
+            'dm-haiku @ https://github.com/deepmind/dm-haiku/archive/v0.0.2.zip'
         ],
         'examples': ['matplotlib', 'seaborn'],
     },


Test patch:
diff --git a/test/contrib/test_module.py b/test/contrib/test_module.py
new file mode 100644
--- /dev/null
+++ b/test/contrib/test_module.py
@@ -0,0 +1,41 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+import flax
+import haiku as hk
+import numpy as np
+
+import numpyro
+from numpyro import handlers
+from numpyro.contrib.module import flax_module, haiku_module
+
+X = np.arange(100)
+Y = 2 * X + 2
+
+
+def haiku_model(x, y):
+    linear_module = hk.transform(lambda x: hk.Linear(100)(x))
+    nn = haiku_module("nn", linear_module, (100,))
+    mean = nn(x)
+    numpyro.sample("y", numpyro.distributions.Normal(mean, 0.1), obs=y)
+
+
+def flax_model(x, y):
+    linear_module = flax.nn.Dense.partial(features=100)
+    nn = flax_module("nn", linear_module, (100,))
+    mean = nn(x)
+    numpyro.sample("y", numpyro.distributions.Normal(mean, 0.1), obs=y)
+
+
+def test_flax_module():
+    with handlers.trace() as flax_tr, handlers.seed(rng_seed=1):
+        flax_model(X, Y)
+    assert flax_tr["nn$params"]['value']['kernel'].shape == (100, 100)
+    assert flax_tr["nn$params"]['value']['bias'].shape == (100,)
+
+
+def test_haiku_module():
+    with handlers.trace() as haiku_tr, handlers.seed(rng_seed=1):
+        haiku_model(X, Y)
+    assert haiku_tr["nn$params"]['value']['linear']['w'].shape == (100, 100)
+    assert haiku_tr["nn$params"]['value']['linear']['b'].shape == (100,)