Compute the KL divergence between two distributions

Yes, after sampling, you pass each slice of the trace to the model logp:

logp = model.logp
for point in trace:
    lp.append(logp(point))

Note that prior samples are in a dict of arrays, so you need to slice it manually.