pymc-devs/pymc

Add rewrite for Mixture when `comp_dists` can be "fused"

Open

#6,803 opened on Jun 29, 2023

View on GitHub
 (0 comments) (1 reaction) (0 assignees)Python (7,926 stars) (1,902 forks)batch import
feature requesthelp wantedlogprob

Description

Description

The following distributions are equivalent:

import pymc as pm

pm.Mixture.dist(w=[0.5, 0.5], comp_dists=[pm.Normal.dist(-1), pm.Normal.dist(1)])
pm.Mixture.dist(w=[0.5, 0.5], comp_dists=pm.Normal.dist([-1, 1]))

But the second one is more efficient, because the logp is vectorized among a single batched Normal.

We could add a rewrite in the logprob_rewrites to convert the former to the second, so that users are not penalized from using the first form (which may be more intuitive for some).

Actually that sort of rewrite stack([rv1, rv2]) -> rv3 could be useful in many places in the logprob submodule

Contributor guide

Tech stack
python
Domain
backend
Issue type
feature
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-3 hours
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
fresh
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
clear
Prerequisites
Familiarity with PyMC's Mixture distributionUnderstanding of logp rewrites in PyMC
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
30
Research direction
Investigate the logprob rewrites module in the PyMC codebase (likely in pymc/logprob/rewrites.py or similar). The goal is to add a rewrite that detects when a Mixture's comp dists are all instances of the same distribution (e.g., Normal) with different parameters, and replaces them with a single batched distribution. Look at existing rewrites for patterns. Implement the transformation and add corresponding tests to ensure correctness. Check if any linked PRs or comments provide guidance; currently none.
Add rewrite for Mixture when `comp_dists` can be "fused" · pymc-devs/pymc#6803 | Good First Issue