feat(gh-2142): add VSBC diagnostic

[michaelellis003]

Summary

Implements the VSBC diagnostic requested in #2142.

This follows Yao et al. (2018), Section 3 / Algorithm 2: simulate datasets from the prior predictive, fit a variational approximation for each dataset, estimate marginal calibration probabilities, and test symmetry of p vs 1 - p with a two-sample KS test.

Changes

• add vsbc_diagnostic in numpyro/infer/calibration.py
• export vsbc_diagnostic from numpyro.infer
• add end-to-end tests in test/infer/test_calibration.py
• support sequential / vectorized / parallel dispatch
• validate observed-site plumbing and simulated_data_to_args
• reject discrete latent sample sites explicitly

References

• Yao, Y., Vehtari, A., Simpson, D., and Gelman, A. (2018). Yes, but Did It Work?: Evaluating Variational Inference.
• Talts, S., Betancourt, M., Simpson, D., Vehtari, A., and Gelman, A. (2018). Validating Bayesian Inference Algorithms with Simulation-Based Calibration.

Fixes #2142

[fehiepsi]

LGTM thanks @michaelellis003!

[fehiepsi]

maybe not expose vsbc_diagnostic here, just add calibration here?

[michaelellis003]

Makes sense, fixed! Thanks for the review!

[juanitorduz]

@michaelellis003 can you please rebase/merge with master as the fixes have been merged upstream :) ?