[Feature Request] Show params in render_model()

[karm-patel]

Hi @fehiepsi and team,

Prof. @nipunbatra explained the issue (in pyro) of rendering params in pyro.render_model() and then I added this feature in pyro (PR). I would like to add the same feature in numpyro.

I have tried to add this feature for numpyro and tested it on some examples. Consider the following model,

def model(data):
    m = numpyro.param("m", jnp.array(0))
    sd = numpyro.param("sd", jnp.array([1]),constraint=constraints.positive)
    lambd = numpyro.sample("lambda",dist.Normal(m,sd))
    with numpyro.plate("N", len(data)):
        numpyro.sample("obs", dist.Exponential(lambd), obs=data)

I have added sample_params and params_constraint keys in dictionary returned by contrib.render.get_model_relations(). I added Provenance Tracking for params in get_model_relations to add values in sample_param.

numpyro.contrib.render.get_model_relations(model,model_args=(data,))

data = jnp.ones(10)
numpyro.render_model(model, model_args=(data,), render_distributions=True))

I have added optional argument render_params in render_model() to show params in the plot. Also, render_distributions=True shows the constraints of params with distributions of the sample.

numpyro.render_model(model, model_args=(data,), render_params=True, render_distributions=True)

@fehiepsi, Can you please review this? if it meets your expectaion, then I would like to create a PR for the same.

[fehiepsi]

Thanks @karm-patel! The enhancement looks great to me.