You can use mask at site “Y” as in this thread (need numpyro master branch) or do Bayesian imputation as in this tutorial.
master