Easily repeating distribution batch members?

[Bryor Snefjella]

Hi all,

I'm trying to figure out if there is an implemented, general way of taking a tfp.distribution, repeating certain batch members, and making a new distribution to sample from. The code below works, but gets cumbersome quickly when I need to do the repetition on many different types of distributions.

import tensorflow as tf

import tensorflow_probability as tfp

norm = tfp.distributions.Normal([[0,1],[2,3]], [[1,2],[3,4]])

### what if you want to repeat the first distribution multiple times? norm_repeat = tfp.distributions.Normal(tf.gather(norm.loc, [0,0,0,1]), tf.gather(norm.scale, [0,0,0,1]))

norm_repeat.sample()

Best,

Bryor Snefjella