How to return documents assigned to topics?

Hi, I’m trying LDA tutorial Automatic autoencoding variational Bayes for latent dirichlet allocation with PyMC3 — PyMC3 3.10.0 documentation, how could I return not only topics with assigned words, but also documents passed to a model with assigned discovered topics?
So it would be not:

but for each topic list all documents assigned to it

like in gensim it’s just lda_model[corpus[0]] which provides with e.g. [(10, 0.998)] where first document is 99.8% about topic 10
what is the analogous for that in pymc3?