pymc-devs / pymc-devs/pymc

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

Open
#6,803 0 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

feature request help wanted logprob
Dominant language
Python
Stars
9.8k
Forks
2.3k
Avg merge
21h 39m
Merged PRs (30d)
5

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in the logprob_rewrites and logprob submodule, using the two Mixture.dist examples as the behavioral reference. Determine how a list of component distributions can be recognized as equivalent to one batched distribution. Done means the separate-component form receives the same vectorized logp behavior as the fused form, with coverage for the demonstrated Normal case.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.