Inference for gaussian_hmm is broken on latest jax version (0.4.30)

[damonbayer]

The following works with jax/jaxlib 0.4.29, but is broken with jax/jaxlib 0.4.30:

Copying from the current documentation:

import numpy as np
import numpyro
import numpyro.distributions as dist
from numpyro.contrib.control_flow import scan

def gaussian_hmm(y=None, T=10):
    def transition(x_prev, y_curr):
        x_curr = numpyro.sample('x', dist.Normal(x_prev, 1))
        y_curr = numpyro.sample('y', dist.Normal(x_curr, 1), obs=y_curr)
        return x_curr, (x_curr, y_curr)

    x0 = numpyro.sample('x_0', dist.Normal(0, 1))
    _, (x, y) = scan(transition, x0, y, length=T)
    return (x, y)


with numpyro.handlers.seed(rng_seed=0):  # generative
    x, y = gaussian_hmm()

If I then try to perform inference:

from jax import random
from numpyro.infer import MCMC, NUTS
nuts_kernel = NUTS(gaussian_hmm)

mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=1000)
rng_key = random.PRNGKey(0)

mcmc.run(rng_key, y=y)

I get the error

TypeError: body_fun output and input must have identical types, got
('ShapedArray(int32[], weak_type=True)', ['ShapedArray(float32[10])', 'DIFFERENT ShapedArray(int32[], weak_type=True) vs. ShapedArray(float0[])', 'ShapedArray(float32[])'], []).

[fehiepsi]

Hi @damonbayer, it is fixed in #1817. We will make a patch release this week.