Using scale in numpyro

I have a numpyro model that is pretty much a very simple logistic regression. Each data point in my dataset has its own weight, so I want to weight them accordingly in terms of the log-likelihood, so the posterior will look something like

log p(theta|x) \proporto w*log p(x|theta) + log p(theta)

where w is the weight associated with x. I’m trying to use the scale effect handler to do this, i.e.

model  = scale(model, weights) 

where weights is a numpy array consist of weights for the data points. But apparently this did not work, as the model did not scale the log-likelihood with the weights at all. Am I doing something wrong?

here’s an example:

weights = jnp.array([1.1, 2.2, 3.3])
with numpyro.plate("data", 3), numpyro.handlers.scale(scale=weights):
    numpyro.sample("obs", dist.Bernoulli(logits=jnp.zeros(3)), obs=jnp.ones(3))
2 Likes

That worked. Thanks!