log_prob issue when concatenating distributions across batch_shape

Dear all,

I was wondering if there is any way of batching already built distributions along a certain batch_shape axis. For example:

A = tfd.Normal(np.random.normal(size=[10, 2, 4]), np.random.normal(size=[10, 2, 4]))
B = tfd.Normal(np.random.normal(size=[10, 4, 4]), np.random.normal(size=[10, 4, 4]))

# print(A); print(B)
tfp.distributions.Normal("Normal", batch_shape=[10, 2, 4], event_shape=[], dtype=float64)
tfp.distributions.Normal("Normal", batch_shape=[10, 4, 4], event_shape=[], dtype=float64)

# Concatenate A and B and get ->
<tfp.distributions.Normal 'Normal' batch_shape=[10, 6, 4] event_shape=[] dtype=float64>

I tried using tfd.Blockwise, but the 'concatenated' axis ends up being part of the event_shape. On a similar note, is there any way of 'converting' an event_shape dimension to a batch_shape one? (something like an inverse of tfd.Independent).

l = [A[:, i] for i in range(A.batch_shape[1])] + [B[:, i] for i in range(B.batch_shape[1])]    
tfd.Blockwise(l)

<tfp.distributions.Blockwise 'Blockwise' batch_shape=[10, 4] event_shape=[6] dtype=float64>

Thank you very much!
Lucas

Start by reproducing the A and B tfd.Normal examples and the tfd.Blockwise result shown in the issue. Read the batch_shape and event_shape behavior of tfd.Normal and tfd.Blockwise; done means establishing whether concatenation across a batch axis or conversion from event_shape to batch_shape is supported and documenting the resulting behavior.