How to deal with missing values

ricardoV94 · July 3, 2024, 10:19am

Automatic imputation is compatible with JAX, it’s even tested in our CI:

pymc-devs/pymc/blob/20a69db8f984f16a09aa2055c01da95c30b6a0c2/tests/sampling/test_jax.py#L483-L491


      
          def test_sample_partially_observed():
              with pm.Model() as m:
                  with pytest.warns(ImputationWarning):
                      x = pm.Normal("x", observed=np.array([0, 1, np.nan]))
                  idata = pm.sample(nuts_sampler="numpyro", chains=1, draws=10, tune=10)
          
              assert idata.observed_data["x_observed"].shape == (2,)
              assert idata.posterior["x_unobserved"].shape == (1, 10, 1)
              assert idata.posterior["x"].shape == (1, 10, 3)

You may however have a model based on mutable data / coords which is dynamic in shape by default and which JAX can’t handle. In that case you can:

Pass an explicit shape to pm.Data and your observed variable (can do this alongside the dims)
Call freeze_data_and_dims prior to JAX sampling: pymc.model.transform.optimization.freeze_dims_and_data — PyMC v5.16.1 documentation

You’ll probably need the latest version of PyMC.

The reason it fails is probably specific to your model, hence why I asked for more details

Topic		Replies	Views
Masking missing values of predictors Questions	3	1385	July 10, 2020
Missing values in a model? Questions	12	4791	November 7, 2018
Missing values -- unexpected behavior Questions	0	465	February 19, 2020
Bayesian VAR with missing data: PyTensor error when sampling with blackjax or numpyro v5 bug , modeling	3	431	June 2, 2023
Disabling missing data imputation Questions	17	2261	October 10, 2023

How to deal with missing values

Related topics