JAX/NUMPYRO does not include constant_data

[jbh1128d1]

When I sample larger datasets using jax (which will only sample this without error), the constant_data is not included in the inference_data output. See below.

example.

coords = {
    "obs_id": [0,1,2,3,4],
}
with pm.Model(coords=coords) as rugby_model:
    item_idx = pm.Data('item_idx',[0,1,2,3,4], dims="obs_id", mutable=False)
    a = pm.Normal("a", 0.0, sigma=10.0, shape=5)

    theta = a[item_idx]
    sigma = pm.HalfCauchy("error", 0.5)

    y = pm.Normal("y", theta, sigma=sigma, observed=[3,2,6,8,4])

    idata= pm.sample()
    idata_jax = pymc.sampling_jax.sample_numpyro_nuts(tune=1000, chains = 4, target_accept=0.9)

Complete error traceback

[The complete error output here]

Please provide any additional information below.

Versions and main components

• PyMC/PyMC3 Version: 4.0.0b6
• Aesara/Theano Version: 2.5.1
• Python Version: 3.7.12
• Operating system: Debian via Google Cloud Platform
• How did you install PyMC/PyMC3: pip install pymc --pre (https://github.com/pymc-devs/pymc/wiki/Installation-Guide-(Linux))

[ricardoV94]

Jeesh how many InferenceData groups are there xD? I confirmed that this is still missing in main. Thanks for reporting

[OriolAbril]

@ricardoV94 worst case scenario below 😜

[twiecki]

@ricardoV94 Can you add some more info on what needs to be done here?

[ricardoV94]

@OriolAbril might have a better idea, I am not very familiar with the InferenceData creation from PyMC

[OriolAbril]

The pm.Data variables are included in the generated InferenceData object in the constant_data group, similarty to the variables passed as observed to any variables in the model are added to the observed_data group. I don't know how the jax sampler wrappers work (at any step, not even inferencedata creation) so I don't know what specific steps need to be taken to fix this, and it also is strange to have observed_data but not constant_data 🤔.

[danhphan]

Hi, I think this issue is due to no constant_data is included in az_trace in pymc.sampling_jax.sample_numpyro_nuts(). Similar things also happen in pymc.sampling_jax.sample_blackjax_nuts().

#5189 has done a great job to add log_likelihood, observed_data, and sample_stats to numpyro sampler.

So, we can add constant_data into the sample_numpyro_nuts function by creating a function find_constant_data(model), similar to find_observations(model) for observed_data.

In a couple of days, if no-one take this. I am happy to do this task :)

[ricardoV94]

@danhphan is this one resolved?

[danhphan]

Hi @ricardoV94 , yes, I think it was fixed on this PR #5807