title: A JAX custom VJP function for multiple input variable does not work for NumPyro/HMC-NUTS
body: <p>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:</p>
  <pre class="lang-py prettyprint-override"><code>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)
  </code></pre>
  <p>Here, I defined h(x)=sin(x) by manual. Then, I made a test data as</p>
  <pre class="lang-py prettyprint-override"><code>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)
  </code></pre>
  <p><a href="https://i.stack.imgur.com/3BojZ.png" rel="nofollow noreferrer">test data</a></p>
  <p>I was able to perform a HMC-NUTS in NumPyro in this case as</p>
  <pre class="lang-py prettyprint-override"><code>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()
  </code></pre>
  <p>It works.</p>
  <pre><code>sample: 100%|██████████| 3000/3000 [00:15&lt;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
  </code></pre>
  <p>However, if I define a multi variable function as,</p>
  <pre class="lang-py prettyprint-override"><code>@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)
  </code></pre>
  <p>then perform a HMC-NUTS as</p>
  <pre class="lang-py prettyprint-override"><code>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()
  </code></pre>
  <p>then I got an error as</p>
  <pre><code>TypeError: mul got incompatible shapes for broadcasting: (3,), (22,).
  </code></pre>
  <p>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.</p>