Change scipy.special.logit into jnp.log

[petergchang]

IIUC, the logits argument for jax.random.categorical correspond to log probabilities instead of their logit values
(defined by log(p/(1-p))),
and so the usage of jax.random.categorical in hmm_forwards_filtering_backwards_sampling_jax and hmm_sample_jax should be modified to take in
jnp.log(.) instead of logit(.).

From the documentation for jax.random.catgegorical:

logits - Unnormalized log probabilities of the categorical distribution(s) to sample from, so that softmax(logits, axis) gives the corresponding probabilities.

[murphyk]

Yes, you are right. thanks.