How to access internal probability distribution of TensorFlow/Magenta MusicTransformer model wrapped in tensor2tensor?

Question

I am interested in using Magenta toolbox\'s MusicTransformer model to do some probabilistic analyses of musical corpora. Helpfully, Magenta have released a colab notebook th