HMM.forward_backward returns zero posterior marginals when length is shorter than obs_seq

[satyam-mishra-dev]

HMM.forward_backward documents length as the valid length of an observation sequence, useful for truncating padded sequences. However, when length < len(obs_seq), the returned posterior marginals for valid timesteps can become all zeros.

Minimal repro:

import distrax
import jax.numpy as jnp
from distrax._src.utils import hmm

model = hmm.HMM(
    init_dist=distrax.Categorical(probs=jnp.array([0.6, 0.4])),
    trans_dist=distrax.Categorical(
        probs=jnp.array([[0.8, 0.2], [0.3, 0.7]])
    ),
    obs_dist=distrax.Normal(
        loc=jnp.array([0.0, 3.0]),
        scale=jnp.array([0.5, 0.5]),
    ),
)

obs = jnp.array([0.05, 2.9, 0.1, 99.0, 99.0])

_, beta_pad, gamma_pad, ll_pad = model.forward_backward(obs, length=jnp.array(3))
_, beta_prefix, gamma_prefix, ll_prefix = model.forward_backward(obs[:3])

print(ll_pad, ll_prefix)
print(beta_pad[:3])
print(beta_prefix)
print(gamma_pad[:3])
print(gamma_prefix)

Observed:

ll_pad == ll_prefix

beta_pad[:3]
[[0. 0.]
 [0. 0.]
 [0. 0.]]

gamma_pad[:3]
[[0. 0.]
 [0. 0.]
 [0. 0.]]

Expected:
forward_backward(obs, length=3) should match forward_backward(obs[:3]) for the first three valid timesteps.

Likely cause:
In HMM.backward, padded suffix steps replace the backward carry with zeros when t > length. Those zeros then propagate backward into valid timesteps.

Possible fix:
For padded suffix steps, preserve beta_prev instead of replacing it with zeros, and add a regression test comparing forward_backward(obs, length=L) against forward_backward(obs[:L]).