Repository: pyro-ppl/numpyro

Issue:
More than 1 `input_shape` when initializing `flax_module`
Some modules require more than 1 input when initializing, which can be passed through `kwargs`. But this doesn't work in some cases. For example:

```python
class RNN(nn.Module):
    @functools.partial(
        nn.transforms.scan,
        variable_broadcast='params',
        split_rngs={'params': False})
    @nn.compact
    def __call__(self, state, x):
        return RNNCell()(state, x)
``` 

I tried to declare this with the following statement:

```python
rnn = flax_module(
    'rnn',
    RNN(),
    input_shape=(num_hiddens,),
    x=jnp.ones((10, 10))
)
```

But I can't use kwargs because `nn.transforms.scan` does not support them:
```
RuntimeWarning: kwargs are not supported in scan, so "x" is(are) ignored
```

I worked around this by wrapping my `RNN` with another class, after which I could pass `x` as a kwarg. However, I think `input_shape` should allow passing dimensions for more than one input.

https://github.com/pyro-ppl/numpyro/blob/0bff074a4a54a593a7fab7e68b5c10f85dd332a6/numpyro/contrib/module.py#L83


Discussion:
Hi @UmarJ, I think you can ignore that argument and use `state=..., x=...` directly (docs of kwargs: "optional keyword arguments to initialize flax neural network as an alternative to `input_shape`")
> Hi @UmarJ, I think you can ignore that argument and use `state=..., x=...` directly (docs of kwargs: "optional keyword arguments to initialize flax neural network as an alternative to `input_shape`")

Like I said, kwargs don't work when `__call__` is wrapped in `nn.transforms.scan`, for example the RNN model above.
```
RuntimeWarning: kwargs are not supported in scan, so "x" is(are) ignored
```

It's a minor issue, just thought I'd raise it in case someone else was having a similar problem.

Thanks for explaining! I misinterpret the issue.

I think input_shape is a bad pattern here. Originally we would like to keep the same signature w.r.t. jax stax. Maybe it is time to support `*args`? (I made the [wrong comment](https://github.com/pyro-ppl/numpyro/pull/822#discussion_r528119575) here against `args` - seems like it is useful for this usage case.

Resolving patch:
diff --git a/numpyro/contrib/module.py b/numpyro/contrib/module.py
--- a/numpyro/contrib/module.py
+++ b/numpyro/contrib/module.py
@@ -22,7 +22,7 @@
 
 
 def flax_module(
-    name, nn_module, *, input_shape=None, apply_rng=None, mutable=None, **kwargs
+    name, nn_module, *args, input_shape=None, apply_rng=None, mutable=None, **kwargs
 ):
     """
     Declare a :mod:`~flax` style neural network inside a
@@ -44,6 +44,8 @@ def flax_module(
 
     :param str name: name of the module to be registered.
     :param flax.linen.Module nn_module: a `flax` Module which has .init and .apply methods
+    :param args: optional arguments to initialize flax neural network
+        as an alternative to `input_shape`
     :param tuple input_shape: shape of the input taken by the
         neural network.
     :param list apply_rng: A list to indicate which extra rng _kinds_ are needed for
@@ -80,7 +82,7 @@ def flax_module(
 
     if nn_params is None:
         # feed in dummy data to init params
-        args = (jnp.ones(input_shape),) if input_shape is not None else ()
+        args = (jnp.ones(input_shape),) if input_shape is not None else args
         rng_key = numpyro.prng_key()
         # split rng_key into a dict of rng_kind: rng_key
         rngs = {}
@@ -121,7 +123,7 @@ def apply_without_state(params, *args, **kwargs):
     return partial(apply_fn, nn_params)
 
 
-def haiku_module(name, nn_module, *, input_shape=None, apply_rng=False, **kwargs):
+def haiku_module(name, nn_module, *args, input_shape=None, apply_rng=False, **kwargs):
     """
     Declare a :mod:`~haiku` style neural network inside a
     model so that its parameters are registered for optimization via
@@ -143,6 +145,8 @@ def haiku_module(name, nn_module, *, input_shape=None, apply_rng=False, **kwargs
     :param str name: name of the module to be registered.
     :param nn_module: a `haiku` Module which has .init and .apply methods
     :type nn_module: haiku.Transformed or haiku.TransformedWithState
+    :param args: optional arguments to initialize flax neural network
+        as an alternative to `input_shape`
     :param tuple input_shape: shape of the input taken by the
         neural network.
     :param bool apply_rng: A flag to indicate if the returned callable requires
@@ -178,7 +182,7 @@ def haiku_module(name, nn_module, *, input_shape=None, apply_rng=False, **kwargs
         assert (nn_state is None) == (nn_params is None)
 
     if nn_params is None:
-        args = (jnp.ones(input_shape),) if input_shape is not None else ()
+        args = (jnp.ones(input_shape),) if input_shape is not None else args
         # feed in dummy data to init params
         rng_key = numpyro.prng_key()
         if with_state:
@@ -243,7 +247,14 @@ def _update_params(params, new_params, prior, prefix=""):
 
 
 def random_flax_module(
-    name, nn_module, prior, *, input_shape=None, apply_rng=None, mutable=None, **kwargs
+    name,
+    nn_module,
+    prior,
+    *args,
+    input_shape=None,
+    apply_rng=None,
+    mutable=None,
+    **kwargs
 ):
     """
     A primitive to place a prior over the parameters of the Flax module `nn_module`.
@@ -282,6 +293,8 @@ def __call__(self, x):
             prior={"bias": dist.Cauchy(), "kernel": dist.Normal()}
 
     :type prior: dict, ~numpyro.distributions.Distribution or callable
+    :param args: optional arguments to initialize flax neural network
+        as an alternative to `input_shape`
     :param tuple input_shape: shape of the input taken by the neural network.
     :param list apply_rng: A list to indicate which extra rng _kinds_ are needed for
         ``nn_module``. For example, when ``nn_module`` includes dropout layers, we
@@ -355,6 +368,7 @@ def __call__(self, x):
     nn = flax_module(
         name,
         nn_module,
+        *args,
         input_shape=input_shape,
         apply_rng=apply_rng,
         mutable=mutable,
@@ -369,7 +383,7 @@ def __call__(self, x):
 
 
 def random_haiku_module(
-    name, nn_module, prior, *, input_shape=None, apply_rng=False, **kwargs
+    name, nn_module, prior, *args, input_shape=None, apply_rng=False, **kwargs
 ):
     """
     A primitive to place a prior over the parameters of the Haiku module `nn_module`.
@@ -391,6 +405,8 @@ def random_haiku_module(
             prior={"bias": dist.Cauchy(), "kernel": dist.Normal()}
 
     :type prior: dict, ~numpyro.distributions.Distribution or callable
+    :param args: optional arguments to initialize flax neural network
+        as an alternative to `input_shape`
     :param tuple input_shape: shape of the input taken by the neural network.
     :param bool apply_rng: A flag to indicate if the returned callable requires
         an rng argument (e.g. when ``nn_module`` includes dropout layers). Defaults
@@ -402,7 +418,7 @@ def random_haiku_module(
     :returns: a sampled module
     """
     nn = haiku_module(
-        name, nn_module, input_shape=input_shape, apply_rng=apply_rng, **kwargs
+        name, nn_module, *args, input_shape=input_shape, apply_rng=apply_rng, **kwargs
     )
     params = nn.args[0]
     new_params = deepcopy(params)


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
@@ -156,7 +156,7 @@ def test_update_params():
 
 
 @pytest.mark.parametrize("backend", ["flax", "haiku"])
-@pytest.mark.parametrize("init", ["shape", "kwargs"])
+@pytest.mark.parametrize("init", ["args", "shape", "kwargs"])
 @pytest.mark.parametrize("callable_prior", [True, False])
 def test_random_module_mcmc(backend, init, callable_prior):
 
@@ -185,9 +185,14 @@ def test_random_module_mcmc(backend, init, callable_prior):
     labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1))
 
     if init == "shape":
+        args = ()
         kwargs = {"input_shape": (3,)}
     elif init == "kwargs":
+        args = ()
         kwargs = {kwargs_name: data}
+    elif init == "args":
+        args = (np.ones(3, dtype=np.float32),)
+        kwargs = {}
 
     if callable_prior:
         prior = (
@@ -197,7 +202,7 @@ def test_random_module_mcmc(backend, init, callable_prior):
         prior = {bias_name: dist.Cauchy(), weight_name: dist.Normal()}
 
     def model(data, labels):
-        nn = random_module("nn", linear_module, prior=prior, **kwargs)
+        nn = random_module("nn", linear_module, prior, *args, **kwargs)
         logits = nn(data).squeeze(-1)
         numpyro.sample("y", dist.Bernoulli(logits=logits), obs=labels)