DOC: add out-of-sample predictions subsection to NumPyro conversion guide by Shlokpalrecha · Pull Request #178 · arviz-devs/arviz-base

Shlokpalrecha · 2026-03-21T08:18:13Z

Context

This PR adds the missing out-of-sample predictions example requested in arviz-base issue #2452, following guidance from @OriolAbril.

What this adds

One focused subsection in the existing NumPyro conversion guide notebook.
Reuses the fitted eight schools mcmc object already present in the notebook.
Generates 2-3 random sigma values and builds out-of-sample predictions for y.
Converts with:
- predictions=...
- predictions_constant_data={"sigma": sigma_oos}
Shows where results appear by inspecting:
- idata_mcmc_oos.predictions
- idata_mcmc_oos.predictions_constant_data

Validation

tox -e check passed locally.
All cells executed successfully in notebook kernel.

review-notebook-app · 2026-03-21T08:18:19Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

read-the-docs-community · 2026-03-21T08:19:49Z

Documentation build overview

📚 arviz-base | 🛠️ Build #31997034 | 📁 Comparing ec46b02 against latest (374cd28)

🔍 Preview build

Show files changed (14 files in total): 📝 14 modified | ➕ 0 added | ➖ 0 deleted

File	Status
api/index.html	📝 modified
how_to/ConversionGuideNumPyro.html	📝 modified
_modules/arviz_base/base.html	📝 modified
_modules/arviz_base/converters.html	📝 modified
_modules/arviz_base/rcparams.html	📝 modified
_modules/arviz_base/reorg.html	📝 modified
api/generated/arviz_base.dataset_to_dataarray.html	📝 modified
api/generated/arviz_base.dataset_to_dataframe.html	📝 modified
api/generated/arviz_base.dict_to_dataset.html	📝 modified
api/generated/arviz_base.explode_dataset_dims.html	📝 modified
api/generated/arviz_base.extract.html	📝 modified
api/generated/arviz_base.make_attrs.html	📝 modified
api/generated/arviz_base.references_to_dataset.html	📝 modified
api/generated/arviz_base.xarray_sel_iter.html	📝 modified

kylejcaron · 2026-03-21T11:42:26Z



 def _check_tilde_start(x):
+    """Check whether an item starts with the negation prefix.


Is there any reason why youre making changes to the utils file?

I think it would be best to keep this PR scoped to just the conversion guide

kylejcaron · 2026-03-21T12:21:08Z

@@ -121,12 +121,12 @@
     "name": "stderr",


Line #2. sigma_pool = rng.uniform(8.0, 18.0, size=3)
the two line sampling procedure isnt really necessary, you can just hard code a few sigma values or sample them in 1 line instead of two

Reply via ReviewNB

kylejcaron · 2026-03-21T12:21:08Z

@@ -121,12 +121,12 @@
     "name": "stderr",


Line #3. sigma_oos = rng.choice(sigma_pool, size=J, replace=True)
lets choose a different number of schools to make predictions on than 8 to avoid confusion. maybe just 3?

Reply via ReviewNB

kylejcaron · 2026-03-21T12:22:06Z

@@ -121,12 +121,12 @@
     "name": "stderr",


Line #2. idata_mcmc_oos.predictions_constant_data
you can remove this cell, its covered by the last cell which returns idata_mcmc_oos

Reply via ReviewNB

I removed that redundant cell and now only return idata_mcmc_oos in this section.

Shlokpalrecha · 2026-03-22T06:27:56Z

Thanks, updated. I removed unrelated scope from this PR, simplified the sigma sampling, switched to 3-school out-of-sample predictions to avoid confusion, and removed the redundant display cell

kylejcaron · 2026-03-25T11:39:18Z

Thanks, updated. I removed unrelated scope from this PR, simplified the sigma sampling, switched to 3-school out-of-sample predictions to avoid confusion, and removed the redundant display cell

Im not seeing those changes reflected here - maybe check if you pushed the correct commit?

OriolAbril · 2026-03-26T21:10:36Z

@@ -35,7 +35,7 @@
  },


This reads more like a PR review than a tutorial page. How about something like:

A common scenario is to generate predictions for newly obtained data or data that was excluded from the fit. The InferenceData schema has a predictions group to hold this kind of posterior predictive samples.
This example shows how posterior predictive samples generated for newly obtained values of sigma can be incorporated into the generated DataTree

Reply via ReviewNB

OriolAbril · 2026-03-26T21:10:36Z

@@ -35,7 +35,7 @@
  },


sigma_model = np.random.default_rng(2026).choice(np.array([9.0, 12.0, 16.0]), size=J, replace=True)
I think this is still overly complicated and using model in the variable name is potentially confusing. The other variables called model are the functions that do encode a probabilistic model. This could be something like: new_sigma = np.array([-2.0, 7.0, 17.0])

predictions_oos_full = predictive_oos(random.PRNGKey(2), J=J, sigma=sigma_model)

I have barely used numpyro but doing this here with J=8 and then having to index seems again too complicated, this is a very common usecase. My guess is this can be simplified to ..., J=3, sigma=new_sigma) which in turn will allow deleting lines 6-8 and using the output of predictive_oos directly as the predictions argument. I think this is what @kylejcaron was asking for in the previous review round.

Reply via ReviewNB

Thanks, agreed. I simplified this section accordingly in docs/source/how_to/ConversionGuideNumPyro.ipynb:
renamed to new_sigma,
switched to direct J=3 predictions call,
removed the previous indexing-heavy flow.
I also adjusted the predictive setup so the simplified example runs cleanly while keeping the tutorial code straightforward.

Shlokpalrecha · 2026-03-27T04:16:08Z

Apologies for the earlier confusion with what was visible in the PR. I’ve now pushed the corrected branch state and aligned the section with all review points:
docs-only scope,
clearer tutorial wording,
simplified out-of-sample code path (new_sigma, direct J=3).
Thanks again for the detailed review.

kylejcaron · 2026-03-30T20:14:45Z

@@ -35,7 +35,7 @@
  },


Line #1. sigma_model = np.random.default_rng(2026).choice(np.array([9.0, 12.0, 16.0]), size=J, replace=True)
Theres no reason to do sampling with replacement here, and you end up with size J=8 instead of J=3 which would be better

you can just hardcode sigma to a size 3 array

Reply via ReviewNB

kylejcaron · 2026-03-30T20:14:45Z

@@ -35,7 +35,7 @@
  },


Line #3. predictive_oos = Predictive(eight_schools_model, mcmc.get_samples())
mcmc.get_samples() contains theta from the fitted schools.

Our goal is to estimate theta for 3 new schools. One of the possible correct patterns for numpyro here would be:

post_samples = mcmc.get_samples() post_samples.pop("theta") # remove theta so it can be estimated for new schools predictive_oos = Predictive(eight_schools_model, post_samples)

Reply via ReviewNB

Shlokpalrecha marked this pull request as draft March 21, 2026 08:24

kylejcaron reviewed Mar 21, 2026

View reviewed changes

OriolAbril reviewed Mar 26, 2026

View reviewed changes

DOC: simplify out-of-sample NumPyro tutorial section

ec46b02

Shlokpalrecha force-pushed the docs/numpyro-oos-predictions branch from 655e6f6 to ec46b02 Compare March 27, 2026 04:12

Shlokpalrecha requested a review from OriolAbril March 29, 2026 15:50

kylejcaron reviewed Mar 30, 2026

View reviewed changes



		def _check_tilde_start(x):
		"""Check whether an item starts with the negation prefix.

Uh oh!

Uh oh!

Conversation

Shlokpalrecha commented Mar 21, 2026

Context

What this adds

Validation

Uh oh!

review-notebook-app Bot commented Mar 21, 2026

Uh oh!

read-the-docs-community Bot commented Mar 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

Uh oh!

kylejcaron Mar 21, 2026

Choose a reason for hiding this comment

Uh oh!

kylejcaron Mar 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kylejcaron Mar 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kylejcaron Mar 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Shlokpalrecha Mar 27, 2026

Choose a reason for hiding this comment

Uh oh!

Shlokpalrecha commented Mar 22, 2026

Uh oh!

kylejcaron commented Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

OriolAbril Mar 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

OriolAbril Mar 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Shlokpalrecha Mar 27, 2026

Choose a reason for hiding this comment

Uh oh!

Shlokpalrecha commented Mar 27, 2026

Uh oh!

kylejcaron Mar 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kylejcaron Mar 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

read-the-docs-community Bot commented Mar 21, 2026 •

edited

Loading

kylejcaron Mar 21, 2026 •

edited

Loading

kylejcaron Mar 21, 2026 •

edited

Loading

kylejcaron Mar 21, 2026 •

edited

Loading

kylejcaron commented Mar 25, 2026 •

edited

Loading

OriolAbril Mar 26, 2026 •

edited

Loading

OriolAbril Mar 26, 2026 •

edited

Loading

kylejcaron Mar 30, 2026 •

edited

Loading

kylejcaron Mar 30, 2026 •

edited

Loading