pymc-devs/pymc

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

Open

#6,803 opened on 2023年6月29日

GitHub で見る
 (0 comments) (1 reaction) (0 assignees)Python (1,902 forks)batch import
feature requesthelp wantedlogprob

Repository metrics

Stars
 (7,926 stars)
PR merge metrics
 (平均マージ 11d) (30d で 7 merged PRs)

説明

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

コントリビューターガイド