Fix flaky test_change_point_x64

[GaetanLepage]

Increase num_warmup to make test_change_point_x64 less flaky.

Otherwise it happens to fail as follows:

____________________________ test_change_point_x64 _____________________________
[gw6] darwin -- Python 3.12.5 /nix/store/9pj4rzx5pbynkkxq1srzwjhywmcfxws3-python3-3.12.5/bin/python3.12

    def test_change_point_x64():
        # Ref: https://forum.pyro.ai/t/i-dont-understand-why-nuts-code-is-not-working-bayesian-hackers-mail/696
        num_warmup, num_samples = 500, 3000

        def model(data):
            alpha = 1 / jnp.mean(data.astype(np.float32))
            lambda1 = numpyro.sample("lambda1", dist.Exponential(alpha))
            lambda2 = numpyro.sample("lambda2", dist.Exponential(alpha))
            tau = numpyro.sample("tau", dist.Uniform(0, 1))
            lambda12 = jnp.where(jnp.arange(len(data)) < tau * len(data), lambda1, lambda2)
            numpyro.sample("obs", dist.Poisson(lambda12), obs=data)

        # fmt: off
        count_data = jnp.array([
            13, 24, 8, 24, 7, 35, 14, 11, 15, 11, 22, 22, 11, 57, 11, 19, 29, 6, 19, 12, 22,
            12, 18, 72, 32, 9, 7, 13, 19, 23, 27, 20, 6, 17, 13, 10, 14, 6, 16, 15, 7, 2,
            15, 15, 19, 70, 49, 7, 53, 22, 21, 31, 19, 11, 1, 20, 12, 35, 17, 23, 17, 4, 2,
            31, 30, 13, 27, 0, 39, 37, 5, 14, 13, 22])
        # fmt: on

        kernel = NUTS(model=model)
        mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
        mcmc.run(random.PRNGKey(4), count_data)
        samples = mcmc.get_samples()
        tau_posterior = (samples["tau"] * len(count_data)).astype(jnp.int32)
        tau_values, counts = np.unique(tau_posterior, return_counts=True)
        mode_ind = jnp.argmax(counts)
        mode = tau_values[mode_ind]
>       assert mode == 44
E       assert 69 == 44

test/infer/test_mcmc.py:375: AssertionError

Fixes #1862

cc @martinjankowiak

[martinjankowiak]

thanks @GaetanLepage !

though i have to admit i'm confused as to what nixos is and why you'd be worried about this test failure in the first place : )

[GaetanLepage]

though i have to admit i'm confused as to what nixos is and why you'd be worried about this test failure in the first place : )

True I didn't give much context here !
Nix(OS) is a very innovative and powerful technology to build and distribute packages and systems in a reproducible way.
nixpkgs is a huge (> 100k) package repository incarnated by a collection of build recipes (called derivation) that express how packages should be built.
As numpyro is one of them, we regularly update it and running its test suite is part of our build process.

[GaetanLepage]

@martinjankowiak is it fine to merge ?

[GaetanLepage]

Thanks :)