Ignoring NaNs in a 2D array

Hi. I am modeling a 2D timeseries array with missing values. Right now, I am relying on using a masked distribution, e.g., like_dist.mask(...), to handle them. However, I noticed that it actually imputes missing values during sampling. So, when I run Predictive the missing values are filled in the final output. I think I would like to just keep them as NaNs instead. I can obviously manually set them back to NaNs at the end of my program. But I was wondering if there is a way to tell numpyro to just ignore these missing values by default, and possibly cut down on computation time? Any advice would be appreciated. Thanks!

Hi @sharsenij14 currently we don’t have support for this.

Hi @fehiepsi . Thanks for the past reply. Understood. I also noticed that jax doesn’t have the ma module…