Repository: pymc-devs/pymc

Issue:
Implement model transform to remove minibatching operations from graph
### Description

When using `pm.Minibatch`, the `pm.sample_posterior_predictive` returns predictions with the size of the minibatch instead of the full dataset size. To make predictions on the full dataset requires the previous trace to be passed into a new model with a similar setup. For complicated models, this would add several lines of code to create a new model that is almost identical to the previous model.

This enhancement would make it easier to perform posterior predictive checks when using minibatch.

relates to: https://discourse.pymc.io/t/minibatch-not-working/14061/10

### Example scenario:

```python
import numpy as np
import pymc as pm
import arviz as az
import pytensor.tensor as pt

# generate data
N = 10000
P = 3
rng = np.random.default_rng(88)
X_full = rng.uniform(2, 10, size=(N, 3))
beta = np.array([1.5, 0.2, -0.9])
y_full = np.matmul(X_full, beta) + rng.normal(0, 1, size=(N,))
```


### Before:

```python
# minibatch
X_mb, y_mb = pm.Minibatch(X_full, y_full, batch_size=100)

# original minibatch model
with pm.Model() as model_mb:
    b = pm.Normal("b", mu=0, sigma=3, shape=(P,))
    sigma = pm.HalfCauchy("sigma", 1)
    mu = pm.Deterministic("mu", pt.matmul(X_mb, b))
    likelihood = pm.Normal(
        "likelihood", mu=mu, sigma=sigma, observed=y_mb, total_size=N
    )

    fit_mb = pm.fit(
        n=100000,
        method="advi",
        progressbar=True,
        callbacks=[pm.callbacks.CheckParametersConvergence()],
        random_seed=88,
    )
    idata_mb = fit_mb.sample(500)

    pm.sample_posterior_predictive(idata_mb, extend_inferencedata=True)
    idata_mb.posterior = pm.compute_deterministics(
        idata_mb.posterior, merge_dataset=True
    )

# new but similar model to the original
with pm.Model() as model_preds:
    X = pm.Data("X", X_full)
    y = pm.Data("y", y_full)

    b = pm.Normal("b", mu=0, sigma=3, shape=(P,))
    sigma = pm.HalfCauchy("sigma", 1)
    mu = pm.Deterministic("mu", pt.matmul(X, b))
    likelihood = pm.Normal("likelihood", mu=mu, sigma=sigma, observed=y)
    
with model_preds:    
    pm.set_data({"X": X_full})
    ypreds = pm.sample_posterior_predictive(idata_mb)

print(f"Minibatch: {idata_mb.posterior_predictive.likelihood.sizes}")
print(f"Full Data: {ypreds.posterior_predictive.likelihood.sizes}")

# output
Minibatch: Frozen({'chain': 1, 'draw': 500, 'likelihood_dim_2': 100})
Full Data: Frozen({'chain': 1, 'draw': 500, 'likelihood_dim_2': 10000})
```

Hints from issue comments:
@welcome: ![Welcome Banner](https://raw.githubusercontent.com/pymc-devs/brand/main/welcome-bot/BannerWelcome.jpg)]
:tada: Welcome to _PyMC_! :tada: We're really excited to have your input into the project! :sparkling_heart:
<br>If you haven't done so already, please make sure you check out our [Contributing Guidelines](https://www.pymc.io/projects/docs/en/latest/contributing/index.html) and [Code of Conduct](https://github.com/pymc-devs/pymc/blob/main/CODE_OF_CONDUCT.md).
@ricardoV94: This is the correct behavior. Minibatch is defined as a stochastic slice of the random variable. You can define a model without minibatch and use the old trace to sample a full size dataset with posterior predictive @aphc14: I have edited the original post to describe the issue as an enhancement/feature request rather than a bug. I'm not able to modify labels on my end though.@ricardoV94: > I have edited the original post to describe the issue as an enhancement/feature request rather than a bug. I'm not able to modify labels on my end though.

Thanks I did that. Note that my suggestion of a model transform wouldn't do anything automatically. The API would be something like:

```python
with pm.Model() as minibatch_m:
  ... # define model with minibatch
  idata = pm.sample()
  
with remove_minibatch(minibatch_m) as m:
  pm.sample_posterior_predictive(idata)
```

Solution patch:
diff --git a/pymc/model/transform/basic.py b/pymc/model/transform/basic.py
index fcf42fdf8..877814cd6 100644
--- a/pymc/model/transform/basic.py
+++ b/pymc/model/transform/basic.py
@@ -13,9 +13,11 @@
 #   limitations under the License.
 from collections.abc import Sequence
 
-from pytensor import Variable
+from pytensor import Variable, clone_replace
 from pytensor.graph import ancestors
+from pytensor.graph.fg import FunctionGraph
 
+from pymc.data import MinibatchOp
 from pymc.model.core import Model
 from pymc.model.fgraph import (
     ModelObservedRV,
@@ -58,3 +60,25 @@ def parse_vars(model: Model, vars: ModelVariable | Sequence[ModelVariable]) -> l
     else:
         vars_seq = (vars,)
     return [model[var] if isinstance(var, str) else var for var in vars_seq]
+
+
+def remove_minibatched_nodes(model: Model) -> Model:
+    """Remove all uses of pm.Minibatch in the Model."""
+    fgraph, _ = fgraph_from_model(model)
+
+    replacements = {}
+    for var in fgraph.apply_nodes:
+        if isinstance(var.op, MinibatchOp):
+            for inp, out in zip(var.inputs, var.outputs):
+                replacements[out] = inp
+
+    old_outs, old_coords, old_dim_lengths = fgraph.outputs, fgraph._coords, fgraph._dim_lengths  # type: ignore[attr-defined]
+    # Using `rebuild_strict=False` means all coords, names, and dim information is lost
+    # So we need to restore it from the old fgraph
+    new_outs = clone_replace(old_outs, replacements, rebuild_strict=False)  # type: ignore[arg-type]
+    for old_out, new_out in zip(old_outs, new_outs):
+        new_out.name = old_out.name
+    fgraph = FunctionGraph(outputs=new_outs, clone=False)
+    fgraph._coords = old_coords  # type: ignore[attr-defined]
+    fgraph._dim_lengths = old_dim_lengths  # type: ignore[attr-defined]
+    return model_from_fgraph(fgraph, mutate_fgraph=True)


Test patch:
diff --git a/tests/model/transform/test_basic.py b/tests/model/transform/test_basic.py
index 25bf2324e..856fbf0b2 100644
--- a/tests/model/transform/test_basic.py
+++ b/tests/model/transform/test_basic.py
@@ -11,9 +11,11 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+import numpy as np
+
 import pymc as pm
 
-from pymc.model.transform.basic import prune_vars_detached_from_observed
+from pymc.model.transform.basic import prune_vars_detached_from_observed, remove_minibatched_nodes
 
 
 def test_prune_vars_detached_from_observed():
@@ -30,3 +32,20 @@ def test_prune_vars_detached_from_observed():
     assert set(m.named_vars.keys()) == {"obs_data", "a0", "a1", "a2", "obs", "d0", "d1"}
     pruned_m = prune_vars_detached_from_observed(m)
     assert set(pruned_m.named_vars.keys()) == {"obs_data", "a0", "a1", "a2", "obs"}
+
+
+def test_remove_minibatches():
+    data_size = 100
+    data = np.zeros((data_size,))
+    batch_size = 10
+    with pm.Model(coords={"d": range(5)}) as m1:
+        mb = pm.Minibatch(data, batch_size=batch_size)
+        mu = pm.Normal("mu", dims="d")
+        x = pm.Normal("x")
+        y = pm.Normal("y", x, observed=mb, total_size=100)
+
+    m2 = remove_minibatched_nodes(m1)
+    assert m1.y.shape[0].eval() == batch_size
+    assert m2.y.shape[0].eval() == data_size
+    assert m1.coords == m2.coords
+    assert m1.dim_lengths["d"].eval() == m2.dim_lengths["d"].eval()