Add `compute_deterministics` helper

ricardoV94 commented 3 months ago

Description

After #7206 this utility becomes more relevant.

import pymc as pm

with pm.Model(coords={"group": (0, 2, 4)}) as m:
    mu_raw = pm.Normal("mu_raw", 0, 1, dims="group")
    mu = pm.Deterministic("mu", mu_raw.cumsum(), dims="group")

    trace = pm.sample(var_names=["mu_raw"], chains=2, tune=5 draws=5)

assert "mu" not in trace.posterior

with m:
    trace.posterior = pm.compute_deterministics(trace.posterior, merge_dataset=True)

assert "mu" in trace.posterior

Also refactored the base logic to apply a Dataset to a PointFunc which was needed in a couple of places (here and in pymc-experimental).

This function could plausible be used by sample_prior_predictive and sample_posterior_predictive, but it seems that for backward compatibility these return and accept more diverse kinds of inputs. Let me know if I am wrong.

I am not so sure where to place the compute_deterministics. I put it in the sampling module which is a bit odd perhaps.

Related Issue

[x] Closes #7221

Checklist

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

Type of change

[x] New feature / enhancement
[ ] Bug fix
[ ] Documentation
[ ] Maintenance
[ ] Other (please specify):

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

ricardoV94 commented 3 months ago

CC @tomicapretto

ricardoV94 commented 3 months ago

Tests are failing because it's picking up the new scipy but the old arviz. Should work if we rerun again in a couple of hours/days

ricardoV94 commented 3 months ago

@OriolAbril the RTD preview is failing because it installs old arviz. Should we rerun or you are sure the docs are building correctly?

ricardoV94 commented 3 months ago

JAX test is failing with new scipy, should be fixed soon. Since these changes don't touch any related code I'm okay with merging despite them.

OriolAbril commented 3 months ago

Go ahead, there are still things to fix with step methods so if something were to break I'll open a PR

pymc-devs / pymc