Repository: pyro-ppl/numpyro

Issue:
Batching for custom transforms
Batching is currently not implemented for `custom_transforms` like `xlogy` and `xlog1py` so functions using these primitives will fail with `vmap`.


Discussion:
I guess it just needs a few lines of code. Let me do it. :D 
I tried to use `batching.defbroadcasted`, but haven't been able to get it to work. Here's a test case for you 😆

```python

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))
```
I thought that it is something like
```
batching.defvectorized(xlogy.primitive)
```
Oppss, it does not work :D 
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.

Resolving patch:
diff --git a/examples/ucbadmit.py b/examples/ucbadmit.py
--- a/examples/ucbadmit.py
+++ b/examples/ucbadmit.py
@@ -71,7 +71,7 @@ def run_inference(dept, male, applications, admit, rng, args):
     init_params, potential_fn, constrain_fn = initialize_model(
         rng, glmm, dept, male, applications, admit)
     init_kernel, sample_kernel = hmc(potential_fn, algo='NUTS')
-    hmc_state = init_kernel(init_params, args.num_warmup_steps)
+    hmc_state = init_kernel(init_params, args.num_warmup)
     hmc_states = fori_collect(args.num_samples, sample_kernel, hmc_state,
                               transform=lambda hmc_state: constrain_fn(hmc_state.z))
     return hmc_states
@@ -109,7 +109,7 @@ def main(args):
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='UCBadmit gender discrimination using HMC')
     parser.add_argument('-n', '--num-samples', nargs='?', default=2000, type=int)
-    parser.add_argument('--num-warmup-steps', nargs='?', default=500, type=int)
+    parser.add_argument('--num-warmup', nargs='?', default=500, type=int)
     parser.add_argument('--device', default='cpu', type=str, help='use "cpu" or "gpu".')
     args = parser.parse_args()
     main(args)
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -318,7 +318,27 @@ def xlogy(x, y):
     return lax._safe_mul(x, np.log(y))
 
 
+def _xlogy_batching_rule(batched_args, batch_dims):
+    x, y = batched_args
+    bx, by = batch_dims
+    # promote shapes
+    sx, sy = np.shape(x), np.shape(y)
+    nx = len(sx) + int(bx is None)
+    ny = len(sy) + int(by is None)
+    nd = max(nx, ny)
+    x = np.reshape(x, (1,) * (nd - len(sx)) + sx)
+    y = np.reshape(y, (1,) * (nd - len(sy)) + sy)
+    # correct bx, by due to promoting
+    bx = bx + nd - len(sx) if bx is not None else nd - len(sx) - 1
+    by = by + nd - len(sy) if by is not None else nd - len(sy) - 1
+    # move bx, by to front
+    x = batching.move_dim_to_front(x, bx)
+    y = batching.move_dim_to_front(y, by)
+    return xlogy(x, y), 0
+
+
 ad.defjvp(xlogy.primitive, _xlogy_jvp_lhs, _xlogy_jvp_rhs)
+batching.primitive_batchers[xlogy.primitive] = _xlogy_batching_rule
 
 
 def _xlog1py_jvp_lhs(g, x, y):
@@ -337,6 +357,25 @@ def _xlog1py_jvp_rhs(g, x, y):
     return g * lax._safe_mul(x, np.reciprocal(1 + y))
 
 
+def _xlog1py_batching_rule(batched_args, batch_dims):
+    x, y = batched_args
+    bx, by = batch_dims
+    # promote shapes
+    sx, sy = np.shape(x), np.shape(y)
+    nx = len(sx) + int(bx is None)
+    ny = len(sy) + int(by is None)
+    nd = max(nx, ny)
+    x = np.reshape(x, (1,) * (nd - len(sx)) + sx)
+    y = np.reshape(y, (1,) * (nd - len(sy)) + sy)
+    # correct bx, by due to promoting
+    bx = bx + nd - len(sx) if bx is not None else nd - len(sx) - 1
+    by = by + nd - len(sy) if by is not None else nd - len(sy) - 1
+    # move bx, by to front
+    x = batching.move_dim_to_front(x, bx)
+    y = batching.move_dim_to_front(y, by)
+    return xlog1py(x, y), 0
+
+
 @custom_transforms
 def xlog1py(x, y):
     x, y = _promote_args_like(osp_special.xlog1py, x, y)
@@ -344,6 +383,7 @@ def xlog1py(x, y):
 
 
 ad.defjvp(xlog1py.primitive, _xlog1py_jvp_lhs, _xlog1py_jvp_rhs)
+batching.primitive_batchers[xlog1py.primitive] = _xlog1py_batching_rule
 
 
 def entr(p):
@@ -373,6 +413,7 @@ def cumsum(x):
 
 
 ad.defjvp(cumsum.primitive, lambda g, x: np.cumsum(g, axis=-1))
+batching.defvectorized(cumsum.primitive)
 
 
 @custom_transforms
@@ -383,6 +424,7 @@ def cumprod(x):
 # XXX this implementation does not address the case x=0, hence the result in that case will be nan
 # Ref: https://stackoverflow.com/questions/40916955/how-to-compute-gradient-of-cumprod-safely
 ad.defjvp2(cumprod.primitive, lambda g, ans, x: np.cumsum(g / x, axis=-1) * ans)
+batching.defvectorized(cumprod.primitive)
 
 
 def promote_shapes(*args, shape=()):


Test patch:
diff --git a/test/test_distributions_util.py b/test/test_distributions_util.py
--- a/test/test_distributions_util.py
+++ b/test/test_distributions_util.py
@@ -170,6 +170,42 @@ def test_standard_gamma_batch():
         assert_allclose(samples[i], standard_gamma(rngs[i], alphas[i]))
 
 
+@pytest.mark.parametrize('prim', [
+    xlogy,
+    xlog1py,
+])
+def test_binop_batch_rule(prim):
+    bx = np.array([1., 2., 3.])
+    by = np.array([2., 3., 4.])
+    x = np.array(1.)
+    y = np.array(2.)
+
+    actual_bx_by = vmap(lambda x, y: prim(x, y))(bx, by)
+    for i in range(3):
+        assert_allclose(actual_bx_by[i], prim(bx[i], by[i]))
+
+    actual_x_by = vmap(lambda y: prim(x, y))(by)
+    for i in range(3):
+        assert_allclose(actual_x_by[i], prim(x, by[i]))
+
+    actual_bx_y = vmap(lambda x: prim(x, y))(bx)
+    for i in range(3):
+        assert_allclose(actual_bx_y[i], prim(bx[i], y))
+
+
+@pytest.mark.parametrize('prim', [
+    cumsum,
+    cumprod,
+])
+def test_unop_batch_rule(prim):
+    rng = random.PRNGKey(0)
+    bx = random.normal(rng, (3, 5))
+
+    actual = vmap(prim)(bx)
+    for i in range(3):
+        assert_allclose(actual[i], prim(bx[i]))
+
+
 @pytest.mark.parametrize('p, shape', [
     (np.array([0.1, 0.9]), ()),
     (np.array([0.2, 0.8]), (2,)),