Allow prior to be a callable in random module (#1227)

* Solves #1224

Allows prior to be a user defined function in random module

* Changes to #1227

Fixed a mistake in #1227 and added the new funcionality to the docstrings.

* Changes to #1227

Added the requested changes to #1227 and a test for the new functionality, namely `test_random_module_mcmc_callable`.

* Allow prior to be a callable in random module. Fixes PR #1227

Added the changes mentioned by @fehiepsi to PR #1227 in `test_module`.

* Update to PR #1227

Added the changes proposed by @fehiepsi to PR #1227 in order to fix the errors present in it.

* Allow prior to be a callable. Changes to PR #1227

Fixes the errors spotted by @fehiepsi in my previous commits for PR #1227.

* Run black

* Make sure that callable prior is not a distribution

Co-authored-by: Du Phan <[email protected]>

Author: Carlosbogo

numpyro/contrib/module.py

@@ -10,6 +10,7 @@
 from jax.tree_util import register_pytree_node, tree_flatten, tree_unflatten
 
 import numpyro
+import numpyro.distributions as dist
 from numpyro.primitives import mutable as numpyro_mutable
 
 __all__ = [
@@ -223,12 +224,17 @@ def _update_params(params, new_params, prior, prefix=""):
             new_item = new_params[name]
             _update_params(item, new_item, prior, prefix=flatten_name)
         elif (not isinstance(prior, dict)) or flatten_name in prior:
-            d = prior[flatten_name] if isinstance(prior, dict) else prior
             if isinstance(params[name], ParamShape):
                 param_shape = params[name].shape
             else:
                 param_shape = jnp.shape(params[name])
                 params[name] = ParamShape(param_shape)
+            if isinstance(prior, dict):
+                d = prior[flatten_name]
+            elif callable(prior) and not isinstance(prior, dist.Distribution):
+                d = prior(flatten_name, param_shape)
+            else:
+                d = prior
             param_batch_shape = param_shape[: len(param_shape) - d.event_dim]
             # XXX: here we set all dimensions of prior to event dimensions.
             new_params[name] = numpyro.sample(
@@ -270,7 +276,12 @@ def __call__(self, x):
                                      prior={"bias": dist.Cauchy(), "kernel": dist.Normal()},
                                      input_shape=(4,))
 
-    :type prior: dict or ~numpyro.distributions.Distribution
+        Alternatively, we can use a callable. For example the following are equivalent::
+
+            prior=(lambda name, shape: dist.Cauchy() if name == "bias" else dist.Normal())
+            prior={"bias": dist.Cauchy(), "kernel": dist.Normal()}
+
+    :type prior: dict, ~numpyro.distributions.Distribution or callable
     :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
@@ -374,7 +385,12 @@ def random_haiku_module(
                                       prior={"linear.b": dist.Cauchy(), "linear.w": dist.Normal()},
                                       input_shape=(4,))
 
-    :type prior: dict or ~numpyro.distributions.Distribution
+        Alternatively, we can use a callable. For example the following are equivalent::
+
+            prior=(lambda name, shape: dist.Cauchy() if name.startswith("b") else dist.Normal())
+            prior={"bias": dist.Cauchy(), "kernel": dist.Normal()}
+
+    :type prior: dict, ~numpyro.distributions.Distribution or callable
     :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

test/contrib/test_module.py

@@ -148,7 +148,8 @@ def test_update_params():
 
 @pytest.mark.parametrize("backend", ["flax", "haiku"])
 @pytest.mark.parametrize("init", ["shape", "kwargs"])
-def test_random_module_mcmc(backend, init):
+@pytest.mark.parametrize("callable_prior", [True, False])
+def test_random_module_mcmc(backend, init, callable_prior):
 
     if backend == "flax":
         import flax
@@ -179,13 +180,15 @@ def test_random_module_mcmc(backend, init):
     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
+    if callable_prior:
+        prior = (
+            lambda name, shape: dist.Cauchy() if name == bias_name else dist.Normal()
         )
+    else:
+        prior = {bias_name: dist.Cauchy(), weight_name: dist.Normal()}
+
+    def model(data, labels):
+        nn = random_module("nn", linear_module, prior=prior, **kwargs)
         logits = nn(data).squeeze(-1)
         numpyro.sample("y", dist.Bernoulli(logits=logits), obs=labels)