Added from_numpyro_svi converter

[kylejcaron]

Background

Closes #2453 and #2452

Adds support for converting numpyro svi conversion to ArviZ. Also includes support for inferring dims automatically from the model. SVI uses sample dims of ("samples",) instead of ("chain", "draw").

A conversion guide for numpyro was also added, covering the following cases

• MCMC with posterior predictive
• SVI with an autoguide and a posterior predictive
• SVI with a custom guide function and a posterior predictive
• MCMC with auto-labelling event dims

---

📚 Documentation preview 📚: https://arviz-base--95.org.readthedocs.build/en/95/

[codecov-commenter]

Codecov Report

❌ Patch coverage is 19.56522% with 37 lines in your changes missing coverage. Please review.
✅ Project coverage is 50.23%. Comparing base (a634b72) to head (2d484e2).

Files with missing lines	Patch %	Lines
src/arviz_base/io_numpyro.py	17.77%	37 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (a634b72) and HEAD (2d484e2). Click for more details.

HEAD has 2 uploads less than BASE

Flag	BASE (a634b72)	HEAD (2d484e2)
	5	3

Additional details and impacted files

@@             Coverage Diff             @@
##             main      #95       +/-   ##
===========================================
- Coverage   72.34%   50.23%   -22.11%     
===========================================
  Files          19       19               
  Lines        1837     1879       +42     
===========================================
- Hits         1329      944      -385     
- Misses        508      935      +427     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[kylejcaron]

I added an MVP implementation of converting numpyro SVI results to ArviZ. There are two things in the PR that I could use feedback on:

1. the az.from_numpyro_svi inputs

Based on conversations from arviz-devs/arviz#2453, we originally felt this pattern was worth pursuing

az.from_numpyro_svi(
   svi,
   svi_result=svi_result, # contains fitted svi parameters
   model_args=model_args, # args from training. not originally anticipated
   model_kwargs=model_kwargs, # kwargs from training. not originally anticipated
   num_samples=num_samples, # posterior samples. not originally anticipated

   prior=...,
   posterior_predictive=...,
   coords=...,
   ...
)

At the time, I didn't quite realize that model_args, model_kwargs, and num_samples would also be required, bringing the number of extra inputs in the function signature related to svi from 2 to 5 (even if some are optional). Is this too much? Should I instead use a pattern like so?

from arviz_base.io_numpyro import SVIPosterior

posterior = SVIPosterior(
   svi,
   svi_result=svi_result, # the out
   model_args=model_args, # args from training
   model_kwargs=model_kwargs, # kwargs from training
   num_samples=num_samples, # posterior samples
)

az.from_numpyro(
   posterior,
   prior=...,
   posterior_predictive=...,
   coords=...,
   ...
)

2. the structure of the tests

Since numpyro is very functional, the tests required a fair amount of branching logic to handle MCMC vs. SVI with autoguides (which are callables) vs. SVI with custom guides (which are just functions). This added some bloat to the unit tests. Open to suggestions if anyone has any better ideas than how I structured them. Currently they are more concise, but they are less read-able due to all of the conditional statements

[OriolAbril]

sorry about the extremely slow review time. I should be able to review way more quickly going forward. Thanks so much for the great PR with all CI green even ❤️

I agree with most decisions regarding if branching, inputs and tests (haven't really used VI extensively so not the most indicate person to weigh in though).

[kylejcaron]

sorry about the extremely slow review time. I should be able to review way more quickly going forward. Thanks so much for the great PR with all CI green even ❤️

I agree with most decisions regarding if branching, inputs and tests (haven't really used VI extensively so not the most indicate person to weigh in though).

No worries at all, thanks for the review! I'm going to take a look at that Nested Sampler PR you linked as well and try to think about how the patterns I choose here could generalize

[kylejcaron]

@juanitorduz if you have time, inviting you to take a look and give feedback! I know you use numpyro, SVI, and ArviZ quite a bit

[juanitorduz]

Wow! This looks great! I always used the from_dict approach by hand 😅

[OriolAbril]

Left a small comment about the notebook but otherwise I think it is ready to merge. We can then continue iterating over the sample_dims workflow