Conjugate transpose crashes model - tensorflow probability

Hi all,

A quick aside, I tried to add a tag for tensorflow probability, but I did not see an option to add tags. I don’t know whether some features are disabled because I am new to this forum.

I am cross posting a question regarding tensorflow probability which has yet to receive an answer. I traced an error to the conjugate transpose operation, but its not clear to me whether this is a bug or something I am doing incorrectly. Here is what I found:

  1. The model runs in all cases if conjugate = False in tf.transpose

  2. The model runs if conjugate = True and only one chain is ran

  3. The model does not run if conjugate = True and multiple chains are ran

Here is a simple example of the problem:


Any help would be greatly appreciated.