pymc-devs/pymc

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

Open

#6,803 创建于 2023年6月29日

在 GitHub 查看
 (0 评论) (1 反应) (0 负责人)Python (1,902 fork)batch import
feature requesthelp wantedlogprob

仓库指标

Star
 (7,926 star)
PR 合并指标
 (平均合并 11天) (30 天内合并 7 个 PR)

描述

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

贡献者指南