MCMC сохраняет образцы промежуточных значений с апостериорным?

Question

Я использую Pyro для вывода параметров некоторого сложного расчета MCMC. Есть ли способ сохранить значения вычислений, выполненных с предполагаемыми переменными?

Краткая версия кода:

      def model(X_data, Y_data):
    #X_data, and Y_data are column vectors
    A = pyro.sample('A', dist.Normal(loc = 100.0, scale = 10.0)).double()
    B = pyro.sample('B', dist.Normal(loc = 25.0, scale = 2.0)).double()
    C = pyro.sample('C', dist.Normal(loc = 10.0, scale = 1.0)).double()
    D = pyro.sample('D', dist.Normal(loc = 400.0, scale = 50.0)).double()
    
    #Calculate Z for every value of X_data, return as column vector
    Z = some_complicated_calulation(A, B, C, D, X_data)

    #Create a distribution around Z
    Y_predicted = dist.Normal(loc = Z, scale = 10.0)

    #Sample from that distribution, and compare to observed data in Y_data
    pyro.sample('obs', Y_predicted, obs = Y_data)

kernel = NUTS(model, max_tree_depth=3, adapt_step_size=True)
posterior = MCMC(kernel, num_samples=200, warmup_steps=20)
posterior.run(X_grid, Y_grid)

Где X_data а также Y_data - это обучающий набор данных, и оба вектора-столбца тензоров факела (т.е. что-то с формой вроде (8,1)).

Когда я запускаю это и звоню posterior.get_samples()он возвращает заднюю часть A, B, C и D, каждая из которых имеет форму (200). Но мне также нужна задняя часть Z, которая должна иметь форму (8 200).

Сэмплер NUTS делает намного больше, чем 200 (или 220) вычислений Z в фоновом режиме, поэтому использование внешнего списка, тензора или словаря, похоже, не работает. Я думаю, что это можно сделать с pyro.param, но мне непонятно, как это сделать.

Моя конечная цель - получить доверительный интервал some_complicated_calculation для тестового набора неизмеренных значений X.