NotificationsYou must be signed in to change notification settings
Fork2.2k
Star9.4k

Add`vectorize_over_posterior` to`pymc.sampling.forward`#7841

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Jump to bottom

Merged

lucianopaz merged 1 commit intopymc-devs:mainfromlucianopaz:vectorize_post

Jul 4, 2025

Merged

Add`vectorize_over_posterior` to`pymc.sampling.forward`#7841

lucianopaz merged 1 commit intopymc-devs:mainfromlucianopaz:vectorize_post

Jul 4, 2025

Conversation

Copy link

Member

lucianopaz commentedJul 3, 2025•
edited by github-actionsbot
Loading

This PR adds the functionpymc.sampling.forward.vectorize_over_posterior.

This function basically vectorizes the computation of a list of output variables over possible values of random variables or deterministics that are stored in anxarray.Dataset. Any extra random variable remaining in the graph will be vectorized as well, either by changing its size explicitly to match the batch ofchain * draw dimensions, or implicitly given its inputs.

Type of change

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

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

lucianopaz requested a review fromricardoV94

July 3, 2025 13:23

lucianopaz added enhancements pytensor samplers model transform labels

Jul 3, 2025

lucianopaz force-pushed thevectorize_post branch fromcade4b6 to88699ecCompare

July 3, 2025 13:35

ricardoV94 reviewed

Jul 3, 2025

View reviewed changes

Copy link

Member

ricardoV94 left a comment•
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others.Learn more.

Looks good, some minor suggestions and complaints. This PR reminds me that I don't think we havevectorize_node implemented for SymbolicRandomVariables. Shoud be easy to implement, at least for those we have a signature (it's how we know how to resize them automatically)

pymc/pymc/distributions/distribution.py

Lines 395 to 414 in3ae5095

	@_change_dist_size.register(SymbolicRandomVariable)
	defchange_symbolic_rv_size(op:SymbolicRandomVariable,rv,new_size,expand)->TensorVariable:
	extended_signature=op.extended_signature
	ifextended_signatureisNone:
	raiseNotImplementedError(
	f"SymbolicRandomVariable{op} without signature requires custom `_change_dist_size` implementation."
	)

	size=op.size_param(rv.owner)
	ifsizeisNone:
	raiseNotImplementedError(
	f"SymbolicRandomVariable{op} without [size] in extended_signature requires custom `_change_dist_size` implementation."
	)

	params=op.dist_params(rv.owner)

	ifexpandandnotrv_size_is_none(size):
	new_size=tuple(new_size)+tuple(size)

	returnop.rebuild_rv(*params,size=new_size)

tests/sampling/test_forward.pyShow resolvedHide resolved

pymc/sampling/forward.py OutdatedShow resolvedHide resolved

tests/sampling/test_forward.py OutdatedShow resolvedHide resolved

Copy link

Member

ricardoV94 commentedJul 3, 2025•
edited
Loading

Looks good, some minor suggestions and complaints. This PR reminds me that I don't think we havevectorize_node implemented for SymbolicRandomVariables. Shoud be easy to implement, at least for those we have a signature (it's how we know how to resize them automatically)

I'm not suggesting we should do it in this PR, just reminded me

Copy link

codecovbot commentedJul 3, 2025•
edited
Loading

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 92.92%. Comparing base(3ae5095) to head(b2af12f).
Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@##             main    #7841      +/-   ##==========================================+ Coverage   92.91%   92.92%   +0.01%==========================================  Files         107      107                Lines       18286    18313      +27     ==========================================+ Hits        16991    17018      +27  Misses       1295     1295