Shape Issue with Ordered Gaussian Mixture

Questions v5
, ,
1

Hello folks,

I am currently working on a modification of the model from the “Gaussian Mixture Model” example notebook (link).
The original model looks like this.

with pm.Model(coords={"cluster": range(k)}) as model:
    μ = pm.Normal(
        "μ",
        mu=0,
        sigma=5,
        transform=pm.distributions.transforms.ordered,
        initval=[-4, 0, 4],
        dims="cluster",
    )
    σ = pm.HalfNormal("σ", sigma=1, dims="cluster")
    weights = pm.Dirichlet("w", np.ones(k), dims="cluster")
    pm.NormalMixture("x", w=weights, mu=μ, sigma=σ, observed=x)

I am now trying to get another dimension into the model and sample multiple mixtures at once.
Like this:

with pm.Model(coords={"cluster": range(k),"set":range(2)}) as model:
    μ = pm.Normal(
        "μ",
        mu=0,
        sigma=5,
        shape=(2,3),
        transform=pm.distributions.transforms.ordered,
        initval=np.array([[-4, 0, 4],[-4,0,4]]),
        dims=("set","cluster"),
    )
    σ = pm.HalfNormal("σ", sigma=1, dims=("set","cluster"))
    weights = pm.Dirichlet("w", np.ones(k), dims=("set","cluster"))
    pm.NormalMixture("x", w=weights, mu=μ, sigma=σ, observed=np.stack([x,x],axis=-1))

However, with the ordered transformation I am repeatedly running into shape issues, not if I use no transformation. Is this a bug? Does anybody know a way how I can make it work?
Thank you very much in advance

2

Yes, it was reported here `pm.distributions.transforms.ordered` fails on >1 dimension · Issue #5659 · pymc-devs/pymc · GitHub.

There’s a draft PR here Fix `pm.distributions.transforms.ordered` fails on >1 dimension by purna135 · Pull Request #5660 · pymc-devs/pymc · GitHub

If you need it, you could modify your installation with what’s been done in the PR. If you have a look at the PR you’ll see it’s adding a keepdims=True argument in a function call.

2 Likes
3

Thanks for your reply! I’ll have a look.

1 Like