Архитектура сети, используемой в Pyro нормализующих потоках

Question

Я читаю учебник Pyro по нормализации потоков (https://pyro.ai/examples/normalizing_flows_i.html) и хотел бы лучше понять, как эти примеры работают под капотом. Например, я имею в виду архитектуру сети, используемую для получения предельных распределений примера концентрических кругов. В примере базовое распределение (в скрытом пространстве) нормальное, а поток - рациональный сплайн:

      base_dist = dist.Normal(torch.zeros(2), torch.ones(2))
spline_transform = T.Spline(2, count_bins=16)
flow_dist = dist.TransformedDistribution(base_dist, [spline_transform])

Согласно руководству, узлы (сплайна) и их производные - это параметры, которые можно изучить, например, с помощью стохастического градиентного спуска с целью максимального правдоподобия. Далее авторы показывают, как это сделать:

      %%time
steps = 1 if smoke_test else 1001
dataset = torch.tensor(X, dtype=torch.float)
optimizer = torch.optim.Adam(spline_transform.parameters(), lr=1e-2)
for step in range(steps):
    optimizer.zero_grad()
    loss = -flow_dist.log_prob(dataset).mean()
    loss.backward()
    optimizer.step()
    flow_dist.clear_cache()

    if step % 200 == 0:
        print('step: {}, loss: {}'.format(step, loss.item()))

Наконец, автор указывает, как использовать выборку из изученного дистрибутива, чтобы получить новую выборку:

      X_flow = flow_dist.sample(torch.Size([1000,])).detach().numpy()

Я хотел бы знать, какая архитектура NN используется для изучения этих параметров, и есть ли (возможно, простой) способ изменить эту архитектуру.

В более общем плане я хотел бы адаптировать эти простые примеры к одномерному случаю изучения плотности данных временных рядов.