Why inverse log determinant jacobian will be broadcastrd in tfd.TransformedDistribution?

[MokkeMeguru]

In previous, I Implemented Glow in this library. #545 But my training, range of loss is too large.

I try these code. https://colab.research.google.com/gist/MokkeMeguru/650a3cc41769ae8fab9be73f9826644f/matveclu.ipynb

I notice ildj is broadcasted in tfd.TransformedDistribution.log_prob.

I think there are two problems.

1. Blockwise Bijector is invalid implementation.
   folding batch_size is unexpected

2. tfd.TransfromedDistribution._finish_log_prob_for_one_fiber needs some assertion I don't think ildj can be broadcast.

[srvasude]

@csuter / @jvdillon any comments?