Add model_to_minibatch transformation to convert all pm.Data to pm.Minibatch

[jessegrabowski]

Description

A pain point for me when testing different algorithms (e.g. MCMC vs VI) is that I don't want to write a 2nd version of the model with pm.Minibatch on the data.

This PR adds a model transformation that does that for the user. It's the reverse of the remove_minibatched_nodes transformer that @zaxtax implemented recently.

This is a WIP, it doesn't actually work now, because I can't figure out how to rebuild the observed variable with the total_size set correctly. Help wanted.

Related Issue

• Closes #
• Related to #

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

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

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

[ricardoV94]

This is a WIP, it doesn't actually work now, because I can't figure out how to rebuild the observed variable with the total_size set correctly. Help wanted.

You can use the lower level utility:

pymc/pymc/variational/minibatch_rv.py

Line 53 in ef26ae8

def create_minibatch_rv(

Then make that a vanilla observed RV

[ricardoV94]

Ah you already did that, so your question is how to get total size? Grab the batch shape of the variable and constant fold it without raising if it can't be fully folded

[jessegrabowski]

My real issue was not understanding what needs to be the key and value in the replacements, between:

1. The model variable
2. The memo variable
3. The fgraph variable

[ricardoV94]

the best is usual to replace the whole fgraph ModelObservedRV by a new one. You probably have to discard any dims on the batch dimension which is an input to that op

[jessegrabowski]

I don't really understand what that answer means

[ricardoV94]

dprint the fgraph and it will perhaps be more obvious what I am mumbling

[jessegrabowski]

The problem i was running into was that I ended up with two beta RVs after doing the replace. Beta was the only RV implicated in the ModelObservedRV sub-graph

[zaxtax]

Because Minibatch assumes the data variables have the same length, it might make sense to take a variables argument. Or have some way to group data variables of the same size (same dim name maybe?)

…

On Thu, 15 May 2025, 15:35 Ricardo Vieira, ***@***.***> wrote: *ricardoV94* left a comment (pymc-devs/pymc#7785) <#7785 (comment)> dprint the fgraph and it will perhaps be more obvious what I am mumbling — Reply to this email directly, view it on GitHub <#7785 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUMC5VCN6VAAJKNHEMT26SJPZAVCNFSM6AAAAAB5F7LYYKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQOBTHAZTINZXG4> . You are receiving this because you were mentioned.Message ID: ***@***.***>