Add vectorize_over_posterior to pymc.sampling.forward

[lucianopaz]

This PR adds the function pymc.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 an xarray.Dataset. Any extra random variable remaining in the graph will be vectorized as well, either by changing its size explicitly to match the batch of chain * 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/

[ricardoV94]

Looks good, some minor suggestions and complaints. This PR reminds me that I don't think we have vectorize_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 in 3ae5095

	@_change_dist_size.register(SymbolicRandomVariable)
	def change_symbolic_rv_size(op: SymbolicRandomVariable, rv, new_size, expand) -> TensorVariable:
	extended_signature = op.extended_signature
	if extended_signature is None:
	raise NotImplementedError(
	f"SymbolicRandomVariable {op} without signature requires custom `_change_dist_size` implementation."
	)
	size = op.size_param(rv.owner)
	if size is None:
	raise NotImplementedError(
	f"SymbolicRandomVariable {op} without [size] in extended_signature requires custom `_change_dist_size` implementation."
	)
	params = op.dist_params(rv.owner)
	if expand and not rv_size_is_none(size):
	new_size = tuple(new_size) + tuple(size)
	return op.rebuild_rv(*params, size=new_size)

[ricardoV94]

Looks good, some minor suggestions and complaints. This PR reminds me that I don't think we have vectorize_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

[codecov]

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              

Files with missing lines	Coverage Δ
pymc/sampling/forward.py	96.73% <100.00%> (+0.35%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ricardoV94]

Looks great, failing test is unrelated I've seen it fail sporadically