observed_data missing from InferenceData when data is cast

[ricardoV94]

A user on discourse reported that observed_data is missing from the InferenceData returned from the jax samplers in sampling_jax.py. We should check if that is still the case and fix it.

https://discourse.pymc.io/t/pymc-v4-not-returning-observed-data-in-inferencedata-object/8916

[zaxtax]

I just did a cursory check and observed_data is in the InferenceData object returned. Happy to add a unit test for this, but more interested in reproducing the user's error.

[fbarfi]

Have been using sampling_jax with no problem -- observed_data always in inference data object.

[ricardoV94]

Thanks @zaxtax and @fbarfi

I'll close the issue as it seems it was already solved

[vethno]

I apologize for the delay in uploading the minimal example to PyMC Discourse. I am still having the issue.

https://discourse.pymc.io/t/pymc-v4-not-returning-observed-data-in-inferencedata-object/8916/3

Per the post, it works with some distributions, but when I use a mixture distribution like pm.NegativeBinomial() or pm.ZeroInflatedNegativeBinomial()

[zaxtax]

Yep, I am able to reproduce this. Looking into it now.

[zaxtax]

The problem actually shows up if you use regular pm.sample so it's likely something on the distribution as @OriolAbril suggested.

[zaxtax]

It seems to be related to an interaction between pm.Data and pm.NegativeBinomial

The following contains observed_data:

with pm.Model() as model:
     t = pm.Gamma("t", 1, 1)
     y = pm.NegativeBinomial('y', mu=t, alpha=0.5, observed=5)
     trace = pm.sample(tune=20, draws=40)

While this does not:

with pm.Model() as model:
     c = pm.Data("c", 5)
     t = pm.Gamma("t", 1, 1)
     y = pm.NegativeBinomial('y', mu=t, alpha=0.5, observed=c)
     trace = pm.sample(tune=20, draws=40)

[ricardoV94]

Does that problem with pm.Data, happen with other distributions?

[zaxtax]

@ricardoV94 yep. I have replicated it with other discrete distributions like Bernoulli and Binomial

[ricardoV94]

Does it happen with continuous distributions? Might be that the data is being casted to discrete and that's why the InferenceData converter cannot see it.

[zaxtax]

I've possibly narrowed it down to https://github.com/pymc-devs/pymc/blob/main/pymc/aesaraf.py#L234 in the extract_obs_data function. If that isinstance check is removed it works. Maybe this is something @brandonwillard has insights on?