Dirichlet distribution in 4.0.0b4

I replied on the GitHub issue. I couldn’t reproduce it on Colab. Can you try setting up a fresh environment?