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-NUT