Feature request: per-timestep observation uncertainty for Gaussian HMM / ARHMM (heteroscedastic emissions)

[twheatcroft07]

Hi there,

I’m fitting HMM/ARHMM models to latent trajectories from an upstream encoder that outputs:

• mean vector per timestep
• uncertainty (covariance/variance) per timestep

I’m looking for support for known, time-varying observation uncertainty in emissions, e.g.:

y_t given z_t = k is Normal(mean_{t,k}, Sigma_k + R_t)

where:

• Sigma_k is state-specific emission covariance
• R_t is known per-timestep uncertainty from the encoder

Questions:

1. Does this already exist in Dynamax?
2. If not, is there a recommended extension point (custom emissions / inputs)?
3. Is this feature on the roadmap?

If helpful, I’m happy to prototype a PR with guidance on preferred API shape.