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.   
> converting a traditional NN into a Bayesian NN

Yeah, in Pyro, we first use `sample` primitives to get the value of some parameters. Then using `param` primitive for the remaining parameters. I think we can do the same here but need to work around the nested structure of Flax parameter names. I am looking forward to seeing your PR on that. :)

> Do you want both haiku and flax to integrate with numpyro or just one ?

I am seeing your PR which includes two. Let's do it. Thanks!
@fehiepsi I thought when you extract params from a pytorch module, the dictionary of params is also nested ? How do you overcome that in pyro ? 
In PyTorch, you can call [module.named_parameters()](https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.named_parameters) to get a list of parameters with flatten-using-dot names (e.g. `layer1.linear.weight`). Those parameters (together with their flatten names) will be stored in a global ParamsStore. Then we use `pyro.lift` to convert some parameters to random variables.

I think we can use similar approach here. The primitive `random_module` can be something like this
```python
def random_module(name, nn, prior):
    apply_partial = module(name, nn)
    params = apply_partial.args[0]
    new_params = deepcopy(params)
    # loops over the nested dict `new_params`, get flatten name
    # if it appears in `prior` then we replace its value by a `sample` statement
    # (with name of `sample` is `name` + `$$$` + the flatten name)
    # at the same time, we replace the corresponding value in `params` to `None`
    # to not carry those values around in the trace
    return partial(apply_partial.func, new_params, *apply_partial.args[1:], **apply_partial.keywords)
```

I think this way, the implementation will be much simpler than Pyro implementation. If you need me clarify any point, please let me know. Here I assumed `module` returns a partial apply function, with its first arg is `params`. It would be nice if you modify #717 a bit to get that pattern. :)
@TuanNguyen27 Here is an example for looping over a nested dict and updating corresponding params:
```python
from copy import deepcopy
params = {'a': {'b': {'c': {'d': 1}, 'e': 2}, 'f': 3}}
prior = {'a.b.c.d': 4, 'a.f': 5}
new_params = deepcopy(params)

def update_params(prefix, params, new_params, prior):
    for name, item in params.items():
        flatten_name = ".".join([prefix, name]) if prefix else name
        if isinstance(item, dict):
            assert flatten_name not in prior
            new_item = new_params[name]
            update_params(flatten_name, item, new_item, prior)
        elif not isinstance(prior, dict):
            params[name] = None
            new_params[name] = prior
        elif flatten_name in prior:
            params[name] = None
            # NB: replace prior[flatten_name] by a sample statement from it
            new_params[name] = prior[flatten_name]

update_params('', params, new_params, prior)
assert params == {'a': {'b': {'c': {'d': None}, 'e': 2}, 'f': None}}
assert new_params == {'a': {'b': {'c': {'d': 4}, 'e': 2}, 'f': 5}}
```
@fehiepsi thanks for the detailed code example ! 

I'm afraid my current implementation for haiku doesn't allow me to access params from a partial function 

`return lambda x: nn.apply(nn_params, rng_key, x)` is a `<function haiku_module.<locals>.<lambda>` with no attribute. Printing `nn.__dict__` returns `{}`

It is possible to do this for Flax, since the return result of flax_module is `flax.nn.Model(nn, nn_params)`, which has `params` attribute. 

Wondering if there's a work-around to make haiku behave like flax...

Example notebook: https://colab.research.google.com/gist/TuanNguyen27/085c00fa9af777174060430240648449/params.ipynb
Hmm, I think for `haiku`, you can simply return `partial(nn.apply, nn_params, rng_key)` and for `flax`, you can simply return `partial(partial(flax.nn.Model, nn), nn_params)` or `partial(nn.call, nn_params)`. 
interesting, so if we wrap `.apply` with `partial` we can access `nn_params` later on ? 
Yes, you can see how it works in [this comment](https://github.com/pyro-ppl/numpyro/issues/705#issuecomment-681702467).
@TuanNguyen27 I think we have all ingredients for `random_module`. Do you want to pair-code this feature together? In case you agree, could you setup some time this weekend or next week (I believe 30m or 1h is enough)? I would like to hear more about your experience with `pyvarinf` and your ideas on a good example for `random_module`. :)
Yes, i've been wanting to work on this but I forgot, I will ping you directly to set up a time :) 
I realized this is closed, and already implemented, but I'm curious if it makes sense to have 'sister' instantiating functions for haiku_modules that take `*args, **kwargs` rather than input shape. This would enable more complex haiku modules with multiple parameters. For example, 
```python
def haiku_general_module(name, nn_module, *args, **kwargs):
    module_key = name + '$params'
    nn_params = numpyro.param(module_key)
    if nn_params is None:

        # feed in dummy data to init params
        rng_key = numpyro.prng_key()
        nn_params = nn_module.init(rng_key, *args, **kwargs)

        # haiku init returns an immutable dict
        nn_params = hk.data_structures.to_mutable_dict(nn_params)

        # we cast it to a mutable one to be able to set priors for parameters
        # make sure that nn_params keep the same order after unflatten
        params_flat, tree_def = tree_flatten(nn_params)
        nn_params = tree_unflatten(tree_def, params_flat)
        numpyro.param(module_key, nn_params)

    return ft.partial(nn_module.apply, nn_params, None)
```
This should cleanly translate to enable situations like the following:
```python
dual_enc_f = hk.transform(lambda x_0, x_1: DualEncoder(hidden_dim, latent_dim, activate_f)(x_0, x_1))
encoder_m = haiku_general_module("encoder", dual_enc_f, x_0 = x_0_data, x_1 = x_1_data)
z_loc, z_std = encoder_m(x_0_data, x_1_data)
```
Hi @quattro, please feel free to make a pull request to make `haiku_module` (and `flax_module` too) more flexible. I think you don't need to implement a new function. How about
```python
def haiku_module(name, nn_module, *, input_shape=None, **kwargs):
    ...
    if input_shape is not None:
        args = (jnp.ones(input_shape),) if input_shape is not None else ()
    nn_params = nn_module.init(rng_key, *args, **kwargs)
```
?

Resolving patch:
diff --git a/numpyro/contrib/module.py b/numpyro/contrib/module.py
--- a/numpyro/contrib/module.py
+++ b/numpyro/contrib/module.py
@@ -18,7 +18,7 @@
 ]
 
 
-def flax_module(name, nn_module, *, input_shape=None):
+def flax_module(name, nn_module, *, input_shape=None, **kwargs):
     """
     Declare a :mod:`~flax` style neural network inside a
     model so that its parameters are registered for optimization via
@@ -28,6 +28,8 @@ def flax_module(name, nn_module, *, input_shape=None):
     :param flax.nn.Module nn_module: a `flax` Module which has .init and .apply methods
     :param tuple input_shape: shape of the input taken by the
         neural network.
+    :param kwargs: optional keyword arguments to initialize flax neural network
+        as an alternative to `input_shape`
     :return: a callable with bound parameters that takes an array
         as an input and returns the neural network transformed output
         array.
@@ -42,11 +44,10 @@ def flax_module(name, nn_module, *, input_shape=None):
     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.')
+        args = (jnp.ones(input_shape),) if input_shape is not None else ()
         # feed in dummy data to init params
         rng_key = numpyro.prng_key()
-        _, nn_params = nn_module.init(rng_key, jnp.ones(input_shape))
+        _, nn_params = nn_module.init(rng_key, *args, **kwargs)
         # make sure that nn_params keep the same order after unflatten
         params_flat, tree_def = tree_flatten(nn_params)
         nn_params = tree_unflatten(tree_def, params_flat)
@@ -54,7 +55,7 @@ def flax_module(name, nn_module, *, input_shape=None):
     return partial(nn_module.call, nn_params)
 
 
-def haiku_module(name, nn_module, *, input_shape=None):
+def haiku_module(name, nn_module, *, input_shape=None, **kwargs):
     """
     Declare a :mod:`~haiku` style neural network inside a
     model so that its parameters are registered for optimization via
@@ -64,6 +65,8 @@ def haiku_module(name, nn_module, *, input_shape=None):
     :param haiku.Module nn_module: a `haiku` Module which has .init and .apply methods
     :param tuple input_shape: shape of the input taken by the
         neural network.
+    :param kwargs: optional keyword arguments to initialize flax neural network
+        as an alternative to `input_shape`
     :return: a callable with bound parameters that takes an array
         as an input and returns the neural network transformed output
         array.
@@ -79,11 +82,10 @@ def haiku_module(name, nn_module, *, input_shape=None):
     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.')
+        args = (jnp.ones(input_shape),) if input_shape is not None else ()
         # feed in dummy data to init params
         rng_key = numpyro.prng_key()
-        nn_params = nn_module.init(rng_key, jnp.ones(input_shape))
+        nn_params = nn_module.init(rng_key, *args, **kwargs)
         # haiku init returns an immutable dict
         nn_params = haiku.data_structures.to_mutable_dict(nn_params)
         # we cast it to a mutable one to be able to set priors for parameters
@@ -123,7 +125,7 @@ def _update_params(params, new_params, prior, prefix=''):
             new_params[name] = numpyro.sample(flatten_name, d.expand(param_batch_shape).to_event())
 
 
-def random_flax_module(name, nn_module, prior, *, input_shape=None):
+def random_flax_module(name, nn_module, prior, *, input_shape=None, **kwargs):
     """
     A primitive to place a prior over the parameters of the Flax module `nn_module`.
 
@@ -155,6 +157,8 @@ def apply(self, x):
 
     :type param: dict or ~numpyro.distributions.Distribution
     :param tuple input_shape: shape of the input taken by the neural network.
+    :param kwargs: optional keyword arguments to initialize flax neural network
+        as an alternative to `input_shape`
     :returns: a sampled module
 
     **Example**
@@ -210,7 +214,7 @@ def apply(self, x):
         >>> assert losses[-1] < 3000
         >>> assert np.sqrt(np.mean(np.square(y_test - y_pred))) < 1
     """
-    nn = flax_module(name, nn_module, input_shape=input_shape)
+    nn = flax_module(name, nn_module, input_shape=input_shape, **kwargs)
     params = nn.args[0]
     new_params = deepcopy(params)
     with numpyro.handlers.scope(prefix=name):
@@ -219,7 +223,7 @@ def apply(self, x):
     return nn_new
 
 
-def random_haiku_module(name, nn_module, prior, *, input_shape=None):
+def random_haiku_module(name, nn_module, prior, *, input_shape=None, **kwargs):
     """
     A primitive to place a prior over the parameters of the Haiku module `nn_module`.
 
@@ -237,7 +241,7 @@ def random_haiku_module(name, nn_module, prior, *, input_shape=None):
     :param tuple input_shape: shape of the input taken by the neural network.
     :returns: a sampled module
     """
-    nn = haiku_module(name, nn_module, input_shape=input_shape)
+    nn = haiku_module(name, nn_module, input_shape=input_shape, **kwargs)
     params = nn.args[0]
     new_params = deepcopy(params)
     with numpyro.handlers.scope(prefix=name):


Test patch:
diff --git a/test/contrib/test_module.py b/test/contrib/test_module.py
--- a/test/contrib/test_module.py
+++ b/test/contrib/test_module.py
@@ -23,7 +23,7 @@
 from numpyro.infer import MCMC, NUTS
 
 
-def haiku_model(x, y):
+def haiku_model_by_shape(x, y):
     import haiku as hk
 
     linear_module = hk.transform(lambda x: hk.Linear(100)(x))
@@ -32,7 +32,36 @@ def haiku_model(x, y):
     numpyro.sample("y", numpyro.distributions.Normal(mean, 0.1), obs=y)
 
 
-def flax_model(x, y):
+def haiku_model_by_kwargs_1(x, y):
+    import haiku as hk
+
+    linear_module = hk.transform(lambda x: hk.Linear(100)(x))
+    nn = haiku_module("nn", linear_module, x=x)
+    mean = nn(x)
+    numpyro.sample("y", numpyro.distributions.Normal(mean, 0.1), obs=y)
+
+
+def haiku_model_by_kwargs_2(w, x, y):
+    import haiku as hk
+
+    class TestHaikuModule(hk.Module):
+        def __init__(self, dim: int = 100):
+            super().__init__()
+            self._dim = dim
+            return
+
+        def __call__(self, w, x):
+            l1 = hk.Linear(self._dim, name="w_linear")(w)
+            l2 = hk.Linear(self._dim, name="x_linear")(x)
+            return l1 + l2
+
+    linear_module = hk.transform(lambda w, x: TestHaikuModule(100)(w, x))
+    nn = haiku_module("nn", linear_module, w=w, x=x)
+    mean = nn(w, x)
+    numpyro.sample("y", numpyro.distributions.Normal(mean, 0.1), obs=y)
+
+
+def flax_model_by_shape(x, y):
     import flax
 
     linear_module = flax.nn.Dense.partial(features=100)
@@ -41,25 +70,52 @@ def flax_model(x, y):
     numpyro.sample("y", numpyro.distributions.Normal(mean, 0.1), obs=y)
 
 
+def flax_model_by_kwargs(x, y):
+    import flax
+
+    linear_module = flax.nn.Dense.partial(features=100)
+    nn = flax_module("nn", linear_module, inputs=x)
+    mean = nn(x)
+    numpyro.sample("y", numpyro.distributions.Normal(mean, 0.1), obs=y)
+
+
 def test_flax_module():
-    X = np.arange(100)
+    X = np.arange(100).astype(np.float32)
     Y = 2 * X + 2
 
     with handlers.trace() as flax_tr, handlers.seed(rng_seed=1):
-        flax_model(X, Y)
+        flax_model_by_shape(X, Y)
+    assert flax_tr["nn$params"]['value']['kernel'].shape == (100, 100)
+    assert flax_tr["nn$params"]['value']['bias'].shape == (100,)
+
+    with handlers.trace() as flax_tr, handlers.seed(rng_seed=1):
+        flax_model_by_kwargs(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():
-    X = np.arange(100)
+    W = np.arange(100).astype(np.float32)
+    X = np.arange(100).astype(np.float32)
     Y = 2 * X + 2
 
     with handlers.trace() as haiku_tr, handlers.seed(rng_seed=1):
-        haiku_model(X, Y)
+        haiku_model_by_shape(X, Y)
     assert haiku_tr["nn$params"]['value']['linear']['w'].shape == (100, 100)
     assert haiku_tr["nn$params"]['value']['linear']['b'].shape == (100,)
 
+    with handlers.trace() as haiku_tr, handlers.seed(rng_seed=1):
+        haiku_model_by_kwargs_1(X, Y)
+    assert haiku_tr["nn$params"]['value']['linear']['w'].shape == (100, 100)
+    assert haiku_tr["nn$params"]['value']['linear']['b'].shape == (100,)
+
+    with handlers.trace() as haiku_tr, handlers.seed(rng_seed=1):
+        haiku_model_by_kwargs_2(W, X, Y)
+    assert haiku_tr["nn$params"]['value']['test_haiku_module/w_linear']['w'].shape == (100, 100)
+    assert haiku_tr["nn$params"]['value']['test_haiku_module/w_linear']['b'].shape == (100,)
+    assert haiku_tr["nn$params"]['value']['test_haiku_module/x_linear']['w'].shape == (100, 100)
+    assert haiku_tr["nn$params"]['value']['test_haiku_module/x_linear']['b'].shape == (100,)
+
 
 def test_update_params():
     params = {'a': {'b': {'c': {'d': 1}, 'e': np.array(2)}, 'f': np.ones(4)}}
@@ -73,7 +129,9 @@ def test_update_params():
 
 
 @pytest.mark.parametrize("backend", ["flax", "haiku"])
-def test_random_module_mcmc(backend):
+@pytest.mark.parametrize("init", ["shape", "kwargs"])
+def test_random_module__mcmc(backend, init):
+
     if backend == "flax":
         import flax
 
@@ -81,6 +139,7 @@ def test_random_module_mcmc(backend):
         bias_name = "bias"
         weight_name = "kernel"
         random_module = random_flax_module
+        kwargs_name = "inputs"
     elif backend == "haiku":
         import haiku as hk
 
@@ -88,13 +147,7 @@ def test_random_module_mcmc(backend):
         bias_name = "linear.b"
         weight_name = "linear.w"
         random_module = random_haiku_module
-
-    def model(data, labels):
-        nn = random_module("nn", linear_module,
-                           prior={bias_name: dist.Cauchy(), weight_name: dist.Normal()},
-                           input_shape=(dim,))
-        logits = nn(data).squeeze(-1)
-        numpyro.sample("y", dist.Bernoulli(logits=logits), obs=labels)
+        kwargs_name = "x"
 
     N, dim = 3000, 3
     warmup_steps, num_samples = (1000, 1000)
@@ -103,6 +156,18 @@ def model(data, labels):
     logits = np.sum(true_coefs * data, axis=-1)
     labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1))
 
+    if init == "shape":
+        kwargs = {"input_shape": (3,)}
+    elif init == "kwargs":
+        kwargs = {kwargs_name: data}
+
+    def model(data, labels):
+        nn = random_module("nn", linear_module,
+                           {bias_name: dist.Cauchy(), weight_name: dist.Normal()},
+                           **kwargs)
+        logits = nn(data).squeeze(-1)
+        numpyro.sample("y", dist.Bernoulli(logits=logits), obs=labels)
+
     kernel = NUTS(model=model)
     mcmc = MCMC(kernel, warmup_steps, num_samples, progress_bar=False)
     mcmc.run(random.PRNGKey(2), data, labels)