Fix failing LinRegHMM test

[gileshd]

Summary

This PR changes the parameters of the LinearRegressionHMM test to prevent it failing.

Details

At present the LinearRegressionHMM test_sample_and_fit test is failing. The monotonically_increasing assertion is failing as the log likelihoods quickly become jnp.nan.

I think this is fundamentally a question of the stability of the linear solve. I have made some changes to the inputs which are hopefully more robust to solving.

Changes:

• Replace jnp.ones input with jr.normal.
• Reduce size of hidden state to 3.

There are likely more principled ways of choosing a more robust test case but in the interest of getting the test to pass I think it is helpful to make a simple change quickly.

Additional changes:

• Remove unused datetime import and commented lines.