Use pymc.testing for injected code in notebook runner

[Copilot]

Description

Replaces custom mock_sample implementation with pymc.testing.mock_sample from PyMC's standardized testing module (added in pymc-devs/pymc#7887).

Changes:

• Import pymc.testing.mock_sample as base implementation
• Wrap to add sample_stats group for notebook compatibility (divergence checks)
• Fix deprecation: use .sizes instead of .dims
• Correct diverging dtype: bool instead of int

Before:

def mock_sample(*args, **kwargs):
    random_seed = kwargs.get("random_seed", None)
    model = kwargs.get("model", None)
    samples = 10
    idata = pm.sample_prior_predictive(
        model=model,
        random_seed=random_seed,
        draws=samples,
    )
    idata.add_groups(posterior=idata.prior)
    # ... manual cleanup and sample_stats creation

After:

from pymc.testing import mock_sample as pymc_mock_sample

def mock_sample(*args, **kwargs):
    idata = pymc_mock_sample(*args, **kwargs)
    # Add sample_stats for notebook divergence checks
    if "sample_stats" not in idata:
        n_chains = idata.posterior.sizes["chain"]
        n_draws = idata.posterior.sizes["draw"]
        sample_stats = xr.Dataset({
            "diverging": xr.DataArray(
                np.zeros((n_chains, n_draws), dtype=bool),
                dims=("chain", "draw"),
            )
        })
        idata.add_groups(sample_stats=sample_stats)
    return idata

Reduces code duplication, maintains notebook compatibility, improves maintainability.

Related Issue

• Closes [pre-commit.ci] pre-commit autoupdate #2079

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) using numpydoc format.
• If you are a pro: each commit corresponds to a relevant logical change

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Use pymc.testing for injected code in notebook runner</issue_title>
<issue_description>Following merge of pymc-devs/pymc#7887, the pymc.testing module can be used for this runner

pymc-marketing/scripts/run_notebooks/injected.py

Lines 8 to 36 in 9c429ed

	def mock_sample(*args, **kwargs):
	random_seed = kwargs.get("random_seed", None)
	model = kwargs.get("model", None)
	samples = 10
	idata = pm.sample_prior_predictive(
	model=model,
	random_seed=random_seed,
	draws=samples,
	)
	idata.add_groups(posterior=idata.prior)
	# Create mock sample stats with diverging data
	if "sample_stats" not in idata:
	n_chains = 1
	n_draws = samples
	sample_stats = xr.Dataset(
	{
	"diverging": xr.DataArray(
	np.zeros((n_chains, n_draws), dtype=int),
	dims=("chain", "draw"),
	)
	}
	)
	idata.add_groups(sample_stats=sample_stats)
	del idata.prior
	if "prior_predictive" in idata:
	del idata.prior_predictive
	return idata

</issue_description>

<agent_instructions>Keep the solution simple</agent_instructions>

Comments on the Issue (you are @copilot in this section)

• Fixes Use pymc.testing for injected code in notebook runner #2032

---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.

---

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

[juanitorduz]

@copilot add the implementation of the fi, now there is nothing

[juanitorduz]

@copilot add the implementation of the fi, now there is nothing

[williambdean]

The sample_stats in the mock is not in latest pymc version just yet so this is blocked.

Ref the latest two:
https://github.com/pymc-devs/pymc/releases/tag/v5.26.1
https://github.com/pymc-devs/pymc/releases/tag/v5.26.0

[williambdean]

@copilot , go get PyMC team to make a new release

[juanitorduz]

The sample_stats in the mock is not in latest pymc version just yet so this is blocked.

Ref the latest two: https://github.com/pymc-devs/pymc/releases/tag/v5.26.1 https://github.com/pymc-devs/pymc/releases/tag/v5.26.0

why is not included if the PR was merged? 🤔 Ahg, it was merged 2 weeks ago! Got it!

[juanitorduz]

@williambdean, why are the tests 🟢 ? 🤔

[williambdean]

@williambdean, why are the tests 🟢 ? 🤔

What is being done works. However, with the latest version, it can just be

mock_sample = partial(pymc.testing.mock_sample, sample_stats={"diverging": lambda size: np.zeros(size=size)})

There will just be another iteration on this

[juanitorduz]

so shall we merge this one already @williambdean ? (pymc releases can take some time)