Derive probability for broadcasting operations by ricardoV94 · Pull Request #6808 · pymc-devs/pymc

ricardoV94 · 2023-06-30T09:30:44Z

Related to #6398

TODO:

Cover Second/Alloc which are other froms of broadcasting

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

codecov · 2023-06-30T09:42:30Z

Codecov Report

❌ Patch coverage is 97.91667% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.73%. Comparing base (b3033da) to head (1ea3355).

Files with missing lines	Patch %	Lines
pymc/logprob/shape.py	97.87%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6808      +/-   ##
==========================================
- Coverage   91.73%   91.73%   -0.01%     
==========================================
  Files         128      129       +1     
  Lines       20697    20745      +48     
==========================================
+ Hits        18987    19030      +43     
- Misses       1710     1715       +5

Files with missing lines	Coverage Δ
pymc/logprob/__init__.py	`100.00% <100.00%> (ø)`
pymc/logprob/shape.py	`97.87% <97.87%> (ø)`

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

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

larryshamalama · 2023-07-04T03:43:44Z

I will get to this tomorrow morning, sorry about the delay!

larryshamalama

Great work @ricardoV94! :) A lot of nice abstractions, which, together, are why I have many questions

larryshamalama · 2023-07-04T13:59:22Z

+
+
+@_logprob.register(MeasurableBroadcast)
+def broadcast_logprob(op, values, rv, *shape, **kwargs):


Thinking out loud: could this possibly result in inconsistencies elsewhere? For instance, having Mixture components that have been broadcasted which would render them dependent, if that would be an issue

The index mixture only works for basic RVs still so that's fine.

The switch mixture could actually wrongly broadcast the logp. In fact we should also check for invalid switches that mix support dimensions. The current implementation is only correct for ndim_supp==0!

This is another example of why it's so important to have the meta-info for all the MeasurableOps (#6360).

Once we have the meta-info, the Mixture will unambiguously know what kind of measurable variable it is dealing with. In the case of MeasurableBroadcasting, for example, the ndim_supp will have to be at least as large as the number of broadcasted dims (which means we should collapse that logp dimension instead of leaving it as we were doing now!).

We will also know where those support dims are, so that Mixture can know whether we are sub-selecting across core dims.

Without the meta-info, the only way of knowing ndim_supp is by checking the dimensionality of the value vs the logp. We use this logic in some places already:

pymc/pymc/logprob/transforms.py

Lines 432 to 437 in f67ff8b

if input_logprob.ndim < value.ndim:

# For multivariate variables, the Jacobian is diagonal.

# We can get the right result by summing the last dimensions

# of `transform_elemwise.log_jac_det`

ndim_supp = value.ndim - input_logprob.ndim

jacobian = jacobian.sum(axis=tuple(range(-ndim_supp, 0)))

pymc/pymc/logprob/tensor.py

Lines 185 to 189 in f67ff8b

if len({logp.ndim for logp in logps}) != 1:

raise ValueError(

"Joined logps have different number of dimensions, this can happen when "

"joining univariate and multivariate distributions",

)

Which makes me worry whether the probability of a transformed broadcasted variable may be invalid because the "Jacobian" term is going to be counted multiple times?

You raised a very good point, which makes me wonder to what extent #6797 is correct in general?

For instance, if you scale a 3-vector Dirichlet you shouldn't count the Jacobian 3 times, because one of the entries is redundant.

Do we need to propagate information about over-determined elements in multi-dimensional RVs?

The first part of this answer suggests you count it 3 times indeed: https://stats.stackexchange.com/a/487538

I'm surprised :D

Edit: As seen below, that answer is wrong

This I think says something else and correct? https://upcommons.upc.edu/bitstream/handle/2117/366723/p20-CoDaWork2011.pdf?sequence=1&isAllowed=y

I think these should match:

import pymc as pm import numpy as np x = 0.75 print( pm.logp(pm.Beta.dist(5, 9), x).eval(), pm.logp(pm.Dirichlet.dist([5, 9]), [x, 1-x]).eval(), ) # -3.471576058736023 -3.471576058736023 print( pm.logp(2 * pm.Beta.dist(5, 9), 2 * x).eval(), pm.logp(2 * pm.Dirichlet.dist([5, 9]), 2*np.array([x, 1-x])).eval(), ) # -4.164723239295968 -4.857870419855914 print( pm.logp(2 * pm.Beta.dist(5, 9), 2 * x).eval(), (pm.logp(pm.Dirichlet.dist([5, 9]), ([x, 1-x])) - np.log(2)).eval(), ) # -4.164723239295968 -4.164723239295968

Once we have the meta-info, the Mixture will unambiguously know what kind of measurable variable it is dealing with. In the case of MeasurableBroadcasting, for example, the ndim_supp will have to be at least as large as the number of broadcasted dims (which means we should collapse that logp dimension instead of leaving it as we were doing now!).

This makes sense! Would you say that it's better to wait for #6360?

The first part of this answer suggests you count it 3 times indeed: https://stats.stackexchange.com/a/487538

I'm surprised :D

I'm not sure if I fully follow 😅 Nonetheless, I'm glad that this question raised some interesting concerns

larryshamalama · 2023-07-04T14:12:51Z

+    n_new_dims = len(shape) - rv.ndim
+    assert n_new_dims >= 0
+
+    # Enumerate broadcasted dims


Trying to follow along here, this comment is more for "mental scribbles".

rv = pt.random.normal(size=(3, 1)) x = pt.broadcast_to(rv, (5, 2, 3, 4)) # a bit more than your example above # rv.broadcastable = (False, False, False, False) n_new_dims = 2 # 4 - 2 expanded_dims = (0, 1) value.broadcastable[n_new_dims:] = (False, False) # (3, 4) rv.broadcastable = (False, True) # (3, 1) # condition is True only: if (not v_bcast) and rv_bcast = if (not False) and True # condition is True only if v_bast is False and rv_bcast is True broadcast_dims = (3,) # (0 + 2, 1 + 2) but conditions are (False, True)?

review-notebook-app · 2026-07-14T12:07:58Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

read-the-docs-community · 2026-07-14T12:15:19Z

Documentation build overview

📚 pymc | 🛠️ Build #33594975 | 📁 Comparing 1ea3355 against latest (c870b5c)

🔍 Preview build

15 files changed · ± 15 modified

± Modified

Broadcasting lowers to Alloc; the logp scores the base variable at the representative entries and returns -inf where the value is inconsistent with a broadcast, elementwise over the base's batch dimensions. The expanded and broadcast dimensions are degenerate copies of the base entries, so they are consumed like support dimensions and disappear from the logp. Without meta information about the support axes (pymc-devs#6360), other rewrites would treat the copies as independent entries (e.g., counting the jacobian of a transform once per copy), so the broadcast is only claimed when directly valued. For the same reason the implicit broadcasting of measurable inputs inside elemwise transforms is not attempted.

ricardoV94 added enhancements logprob labels Jun 30, 2023

ricardoV94 requested review from larryshamalama and zaxtax June 30, 2023 09:30

ricardoV94 force-pushed the measurable_broadcast branch from 0ec578c to 3fbb4ae Compare June 30, 2023 09:39

ricardoV94 force-pushed the measurable_broadcast branch 3 times, most recently from 2d84ff1 to 042c9f3 Compare June 30, 2023 17:36

larryshamalama reviewed Jul 4, 2023

View reviewed changes

ricardoV94 marked this pull request as draft July 4, 2023 15:39

ricardoV94 mentioned this pull request Jul 5, 2023

Incorporate RV meta information in intermediate MeasurableVariables #6360

Open

fonnesbeck closed this Jun 6, 2026

ricardoV94 reopened this Jun 6, 2026

ricardoV94 force-pushed the measurable_broadcast branch from 042c9f3 to 26bfbdd Compare July 14, 2026 12:07

ricardoV94 changed the base branch from v5 to main July 14, 2026 12:40

ricardoV94 force-pushed the measurable_broadcast branch 2 times, most recently from 1b5896a to af91d6e Compare July 14, 2026 12:40

ricardoV94 force-pushed the measurable_broadcast branch from af91d6e to 1ea3355 Compare July 15, 2026 07:55

ricardoV94 marked this pull request as ready for review July 15, 2026 07:57



		@_logprob.register(MeasurableBroadcast)
		def broadcast_logprob(op, values, rv, shape, *kwargs):

	if input_logprob.ndim < value.ndim:
	# For multivariate variables, the Jacobian is diagonal.
	# We can get the right result by summing the last dimensions
	# of `transform_elemwise.log_jac_det`
	ndim_supp = value.ndim - input_logprob.ndim
	jacobian = jacobian.sum(axis=tuple(range(-ndim_supp, 0)))

	if len({logp.ndim for logp in logps}) != 1:
	raise ValueError(
	"Joined logps have different number of dimensions, this can happen when "
	"joining univariate and multivariate distributions",
	)

Uh oh!

Conversation

ricardoV94 commented Jun 30, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Jun 30, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

larryshamalama commented Jul 4, 2023

Uh oh!

larryshamalama left a comment

Choose a reason for hiding this comment

Uh oh!

larryshamalama Jul 4, 2023

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jul 4, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jul 4, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jul 4, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jul 4, 2023

Choose a reason for hiding this comment

Uh oh!

larryshamalama Jul 5, 2023

Choose a reason for hiding this comment

Uh oh!

larryshamalama Jul 4, 2023

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

review-notebook-app Bot commented Jul 14, 2026

Uh oh!

read-the-docs-community Bot commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ricardoV94 commented Jun 30, 2023 •

edited

Loading

codecov Bot commented Jun 30, 2023 •

edited

Loading

ricardoV94 Jul 4, 2023 •

edited

Loading

ricardoV94 Jul 4, 2023 •

edited

Loading

ricardoV94 Jul 4, 2023 •

edited

Loading

read-the-docs-community Bot commented Jul 14, 2026 •

edited

Loading