Is it possible to use the transforms.AffineCoupling, which takes in a DenseNN and apply the location-scale transformation on all components of the previous random variable.
E.g., to declare a 2d Gaussian with a learnable mean, I would think that one would do
split_dim = 0
dim = 2
hypernet = DenseNN(split_dim, [10*dim], [dim-split_dim, dim-split_dim])
affine = distributions.transforms.AffineCoupling(split_dim,hypernet)
td = distributions.TransformedDistribution(distributions.Normal(torch.zeros(dim),torch.ones(dim)), [affine])
pyro.register("affine",affine)
However, this only works when split_dim >= 1. Am I missing something, or is there another way to do this (except passing a nn.Linear and re-defining the base distribution every time)?