When you want to share code in the future, the easiest way is to use gist: https://gist.github.com/ or just share a Google colab.
I will try to have a look when I have sometime, but it seems your distribution is still very challenging for NUTS. Do you get divergences or converge warnings?