fixup! Add further type annotations to linreghmm

dynamax/hidden_markov_model/models/linreg_hmm.py

@@ -94,7 +94,7 @@ def distribution(
             self,
             params: ParamsLinearRegressionHMMEmissions,
             state: Union[int, Int[Array, ""]],
-            inputs: Array
+            inputs: Float[Array, " input_dim"]
     ):
         prediction = params.weights[state] @ inputs
         prediction +=  params.biases[state]