Batching for custom transforms about numpyro

[fehiepsi]

I guess it just needs a few lines of code. Let me do it. :D

[neerajprad]

I tried to use batching.defbroadcasted, but haven't been able to get it to work. Here's a test case for you 😆

import jax.numpy as np
from jax import random, vmap
from jax.scipy.special import logsumexp
import numpyro.distributions as dist
from numpyro.hmc_util import initialize_model
from numpyro.handlers import sample, seed, substitute, trace
from numpyro.mcmc import mcmc


def beta_bernoulli(data):
    alpha = np.array([1.1, 1.1])
    beta = np.array([1.1, 1.1])
    p_latent = sample('p_latent', dist.Beta(alpha, beta))
    sample('obs', dist.Bernoulli(p_latent), obs=data)
    return p_latent


num_warmup, num_samples = 500, 1000
true_probs = np.array([0.9, 0.1])
data = dist.Bernoulli(true_probs).sample(random.PRNGKey(1), size=(1000, 2))
init_params, potential_fn, constrain_fn = initialize_model(random.PRNGKey(2), beta_bernoulli, data)
samples = mcmc(num_warmup, num_samples, init_params,
               potential_fn=potential_fn,
               constrain_fn=constrain_fn)


def log_lk(rng, params, model, *args, **kwargs):
    model = substitute(seed(model, rng), params)
    model_trace = trace(model).get_trace(*args, **kwargs)
    obs_node = model_trace['obs']
    return np.sum(obs_node['fn'].log_prob(obs_node['value']))


def expected_log_likelihood(rng, params, model, *args, **kwargs):
    n = list(params.values())[0].shape[0]
    log_lk_fn = vmap(lambda rng, params: log_lk(rng, params, model, *args, **kwargs))
    log_lk_vals = log_lk_fn(random.split(rng, n), params)
    return logsumexp(log_lk_vals) - np.log(n)


print(expected_log_likelihood(random.PRNGKey(2), samples, beta_bernoulli, data))

[fehiepsi]

I thought that it is something like

batching.defvectorized(xlogy.primitive)

[fehiepsi]

Oppss, it does not work :D

[neerajprad]

I think defbroadcasting will also broadcast the parameters for binary operators if needed, since its used for binop, but really anything that makes this test case work is fine by me.