Approximation's sample method uses model contexts by zaxtax · Pull Request #7940 · pymc-devs/pymc

zaxtax · 2025-10-29T01:12:13Z

This change makes it so calling sample on fitted approximations uses the model context or an explicitly provided model. Previously, this had a hard assumption that we only want to sample from the same model as the one we fitted.

This also deprecates self.model.

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

📚 Documentation preview 📚: https://pymc--7940.org.readthedocs.build/en/7940/

codecov · 2025-10-31T01:46:40Z

Codecov Report

❌ Patch coverage is 87.50000% with 2 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (v6@4e28043). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
pymc/variational/approximations.py	0.00%	2 Missing ⚠️

Additional details and impacted files

@@          Coverage Diff          @@
##             v6    #7940   +/-   ##
=====================================
  Coverage      ?   91.71%           
=====================================
  Files         ?      124           
  Lines         ?    19967           
  Branches      ?        0           
=====================================
  Hits          ?    18313           
  Misses        ?     1654           
  Partials      ?        0

Files with missing lines	Coverage Δ
pymc/sampling/mcmc.py	`87.99% <100.00%> (ø)`
pymc/variational/opvi.py	`87.42% <100.00%> (ø)`
pymc/variational/approximations.py	`90.36% <0.00%> (ø)`

🚀 New features to boost your workflow:

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

zaxtax · 2025-10-31T03:16:40Z

This still needs some tests that the Deprecation warnings fire

jessegrabowski · 2025-11-04T14:39:44Z

    alias_names = frozenset(["mf"])

-    @node_property
+    @cached_property


these are all just dictionary lookups (or construction of symbolic variables). I don't think it's necessary to cache them

jessegrabowski · 2025-11-04T14:40:35Z

+        # Ensure start is a 1D array and matches ddim
+        start = np.asarray(start).flatten()
+        if start.size != self.ddim:
+            raise ValueError(
+                f"Start array size mismatch: got {start.size}, expected {self.ddim}. "
+                f"Start shape: {start.shape if hasattr(start, 'shape') else 'unknown'}"
+            )


Move this check to _prepare_start

jessegrabowski · 2025-11-04T14:43:25Z

    datalogp_norm = property(lambda self: self.approx.datalogp_norm)
    logq_norm = property(lambda self: self.approx.logq_norm)
-    model = property(lambda self: self.approx.model)
+    model = property(lambda self: modelcontext(None))


If you never refer to self.model now there's no need to store this

jessegrabowski · 2025-11-04T14:46:17Z

+        # Clear old replacements/ordering before rebuilding
+        self.replacements = collections.OrderedDict()
+        self.ordering = collections.OrderedDict()


Is this necessary? Was there a bug?

jessegrabowski · 2025-11-04T14:51:32Z

    def collect(self, item):
        return [getattr(g, item) for g in self.groups]

+    def _variational_orderings(self, model):


The next few functions are new but don't seem to be related to the refactor. Where did these come from?

Also they all take model as an argument, but you should use modelcontext here as well

These are all used in sample to make the logic easier to follow there. We can probably tweak things to have less new functions but that's what they are for.

jessegrabowski · 2025-11-04T14:56:49Z

+            samples: dict = self.sample_dict_fn(draws, model=model, random_seed=random_seed)
+            spec = self._build_trace_spec(model, samples)
+
+            from collections import OrderedDict


Move to top

read-the-docs-community · 2026-03-20T00:45:33Z

Documentation build overview

📚 pymc | 🛠️ Build #32309261 | 📁 Comparing b86f265 against latest (4b0b6ee)

🔍 Preview build

221 files changed · + 4 added · ± 217 modified

+ Added

± Modified

zaxtax · 2026-03-21T23:37:58Z

@jessegrabowski I've tried to incorporate all requested changes. Anything else we need to get this over the line?

jessegrabowski

dang you weren't kidding when you said you slimmed the PR down. Love it.

I approved but pending tests passing ofc

ricardoV94 · 2026-04-16T08:06:39Z

+    with model:
+        x_test = [5, 6, 9, 12, 15]
+        pm.set_data(new_data={"x": x_test}, coords={"obs_id": list(range(len(x_test)))})
+        y_test = pm.sample_posterior_predictive(trace, predictions=True, progressbar=False)
+
+    assert y_test.predictions["obs"].shape == (1, 500, 5)


wouldn't this test have had the same output before your PR? thought you were interested in changing approx.sample() alongside the model (so show how it changes after you call set_data?

Good catch!

ricardoV94

LGTM, can you confirm the test works as regression? Without the changes it failed

zaxtax · 2026-04-17T03:23:45Z

LGTM, can you confirm the test works as regression? Without the changes it failed

The test wasn't a regression. Now it is. Though this did require making more changes

zaxtax · 2026-04-17T14:19:56Z

@ricardoV94 does this still look alright?

ricardoV94 · 2026-04-17T15:05:22Z

        assert data["mu"].shape == ()


+def test_sample_posterior_predictive_after_set_data():


name is not so great anymore?

ricardoV94 · 2026-04-17T17:29:08Z

Thanks @zaxtax

zaxtax requested a review from ricardoV94 October 31, 2025 01:48

zaxtax requested a review from jessegrabowski November 3, 2025 15:39

jessegrabowski requested changes Nov 4, 2025

View reviewed changes

zaxtax force-pushed the removing_model_field_in_fit branch from b4495b7 to 735d8d1 Compare March 20, 2026 00:26

zaxtax changed the base branch from main to v6 March 21, 2026 21:44

zaxtax force-pushed the removing_model_field_in_fit branch from e287859 to e41bb9f Compare March 21, 2026 21:45

zaxtax force-pushed the removing_model_field_in_fit branch 4 times, most recently from ea840de to a64292c Compare April 16, 2026 02:31

jessegrabowski approved these changes Apr 16, 2026

View reviewed changes

zaxtax force-pushed the removing_model_field_in_fit branch from a64292c to fe359ec Compare April 16, 2026 03:24

ricardoV94 reviewed Apr 16, 2026

View reviewed changes

ricardoV94 approved these changes Apr 16, 2026

View reviewed changes

ricardoV94 reviewed Apr 17, 2026

View reviewed changes

zaxtax added 5 commits April 17, 2026 18:18

Don't store the model on the Approximation

47f83af

Add test

c48cf9a

Change set_data between fitting and sampling the posterior

f7267b1

Rework to always require modelcontext

1f244ef

Rename test

b86f265

zaxtax force-pushed the removing_model_field_in_fit branch from 78b809a to b86f265 Compare April 17, 2026 16:19

ricardoV94 added maintenance VI Variational Inference labels Apr 17, 2026

ricardoV94 merged commit 6182801 into pymc-devs:v6 Apr 17, 2026
42 checks passed

zaxtax mentioned this pull request May 29, 2026

Fix regression where passing a model into fit stopped working #8316

Open

11 tasks

		assert data["mu"].shape == ()


		def test_sample_posterior_predictive_after_set_data():

Conversation

zaxtax commented Oct 29, 2025 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Checklist

Type of change

Uh oh!

codecov Bot commented Oct 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

zaxtax commented Oct 31, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Documentation build overview

Uh oh!

zaxtax commented Mar 21, 2026

Uh oh!

jessegrabowski left a comment • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

zaxtax commented Apr 17, 2026

Uh oh!

zaxtax commented Apr 17, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ricardoV94 commented Apr 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zaxtax commented Oct 29, 2025 •

edited by github-actions Bot

Loading

codecov Bot commented Oct 31, 2025 •

edited

Loading

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

edited

Loading

jessegrabowski left a comment •

edited

Loading