A JAX custom VJP function for multiple input variable does not work for NumPyro/HMC-NUTS

Question

I am trying to use a custom VJP (vector-Jacobian product) function as a model for a HMC-NUTS in numpyro. I was able to make a single variable function that works for HMC-NUTS as follows:

import jax.numpy as jnp
from jax import custom_vjp

@custom_vjp
def h(x):
    return jnp.sin(x)

def h_fwd(x):
    return h(x), jnp.cos(x)

def h_bwd(res, u):
    cos_x  = res 
    return (cos_x * u,)

h.defvjp(h_fwd, h_bwd)

Here, I defined h(x)=sin(x) by manual. Then, I made a test data as

import numpy as np
np.random.seed(32)
sigin=0.3
N=20
x=np.sort(np.random.rand(N))*4*np.pi
data=hv(x)+np.random.normal(0,sigin,size=N)

test data

I was able to perform a HMC-NUTS in NumPyro in this case as

import numpyro
import numpyro.distributions as dist

def model(x,y):
    sigma = numpyro.sample('sigma', dist.Exponential(1.))
    x0 = numpyro.sample('x0', dist.Uniform(-1.,1.))
    #mu=jnp.sin(x-x0)
    #mu=hv(x-x0)
    mu=h(x-x0)
    numpyro.sample('y', dist.Normal(mu, sigma), obs=y)

from jax import random
from numpyro.infer import MCMC, NUTS

rng_key = random.PRNGKey(0)
rng_key, rng_key_ = random.split(rng_key)
num_warmup, num_samples = 1000, 2000
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup, num_samples)
mcmc.run(rng_key_, x=x, y=data)
mcmc.print_summary()

It works.

sample: 100%|██████████| 3000/3000 [00:15<00:00, 193.84it/s, 3 steps of size 7.67e-01. acc. prob=0.92]

                mean       std    median      5.0%     95.0%     n_eff     r_hat
     sigma      0.35      0.06      0.34      0.26      0.45   1178.07      1.00
        x0      0.07      0.11      0.07     -0.11      0.26   1243.73      1.00

Number of divergences: 0

However, if I define a multi variable function as,

@custom_vjp
def h(x,A):
    return A*jnp.sin(x)

def h_fwd(x, A):
    res = (A*jnp.cos(x), jnp.sin(x))
    return h(x,A), res

def h_bwd(res, u):
    A_cos_x, sin_x = res
    return (A_cos_x * u, sin_x * u)

h.defvjp(h_fwd, h_bwd)

then perform a HMC-NUTS as

def model(x,y):
    sigma = numpyro.sample('sigma', dist.Exponential(1.))
    x0 = numpyro.sample('x0', dist.Uniform(-1.,1.))
    A = numpyro.sample('A', dist.Exponential(1.))
    mu=h(x-x0,A)
    numpyro.sample('y', dist.Normal(mu, sigma), obs=y)

rng_key = random.PRNGKey(0)
rng_key, rng_key_ = random.split(rng_key)
num_warmup, num_samples = 1000, 2000
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup, num_samples)
mcmc.run(rng_key_, x=x, y=data)
mcmc.print_summary()

then I got an error as

TypeError: mul got incompatible shapes for broadcasting: (3,), (22,).

I suspect that output shape(s) in my function was wrong. But, I could not figure out what was wrong after various trials changing shapes.

Answer 1

def model(x,y):
sigma = numpyro.sample('sigma', dist.Exponential(1.))
x0 = numpyro.sample('x0', dist.Uniform(-1.,1.))
A = numpyro.sample('A', dist.Exponential(1.))
hv=vmap(h,(0,None),0)
mu=hv(x-x0,A)
numpyro.sample('y', dist.Normal(mu, sigma), obs=y)

vmap solved this problem.