gaussian binary tree inference_gym collider model

[gisilvs]

No description provided.

[gisilvs]

@davmre

[davmre]

It would be more efficient to write each layer as a single distribution with batch shape:

layer = yield Root(
  tfd.Normal(loc=initial_loc * tf.ones([2 ** num_layers]),
             scale=initial_scale,
             name='layer_{}'.format(num_layers)))
for l in range(num_layers - 1, 0, -1):
  layer = coupling_link(layer) if coupling_link else layer
  layer = yield tfd.Normal(loc=layer[..., : -1 : 2] - layer[..., 1 : : 2],
                           scale=nodes_scale,
                           name='layer_{}'.format(l))

We'd need to be sure that the CF code does the right thing on batched distributions (which should be treated equivalently to a list of independent dists), but we'd need to do that anyway.