TransformedDistribution.log_prob return shape incompatible with batch dimension semantic

[feynmanliang]

When I run

import simplex.bijectors
import simplex.params

d, params = simplex.bijectors.AffineAutoregressive(
    simplex.params.DenseAutoregressive())(dist.Normal(0, 1))
d.log_prob(torch.zeros((10,1))).shape

I expect the shape of the log_prob of a (sample_shape=[10], batch_shape=[], event_shape=[1]) sample+distribution to have shape [10].

Actual behavior: Result is of shape [10, 10]

RC: TransformedDistribution.log_prob is broadcasting a summation over a row / column vector

[stefanwebb]

Thanks for catching this! I think unit tests for correct output shapes are first on the table :)