search

probml / dynamax

State Space Models library in JAX
https://probml.github.io/dynamax/
MIT License
667 stars 76 forks source link

small edits to initialize? #302

Closed mcusi closed 1 year ago

mcusi commented 1 year ago

Hi, I think I might've noticed a couple small things (although I'm just getting familiar with your repo)

In initialize for DiagonalGaussianHMM, the inputs to NormalInverseGamma are in this order https://github.com/probml/dynamax/blob/71ea2d44860a377409c39706f37a7c551d5f3212/dynamax/hidden_markov_model/models/gaussian_hmm.py#L177 In log_prior these are switched (which I think is correct) https://github.com/probml/dynamax/blob/71ea2d44860a377409c39706f37a7c551d5f3212/dynamax/hidden_markov_model/models/gaussian_hmm.py#L200

the other thing is that I couldn't initialize a Bernoulli HMM with emission_probs because the asserts force them to zero. I think this line should have a 1: https://github.com/probml/dynamax/blob/71ea2d44860a377409c39706f37a7c551d5f3212/dynamax/hidden_markov_model/models/bernoulli_hmm.py#L73

gileshd commented 1 year ago

Hi @mcusi, I think you're right, well noticed and thanks for the heads up! I'll push a fix for these shortly.