text: Title: Using dirichlet sampler directly in Dirichlet distribution
  Body: After https://github.com/google/jax/pull/9906, `jax.random.dirichlet` should be robust for small concentration, so we can remove the current trick that we put in the Dirichlet sampler.