GPU OOM with distributions.Mixture

Hi,

I have an issue with distributions.mixture of 2 normal distributions. For example the little code below take more than 10 Gb of RAM. I don't understand why:

mu = tf.placeholder(tf.float32, (None, 6, 12, 12, 2), name='mu')
sigma = tf.placeholder(tf.float32, (None, 6, 12, 12, 2), name='sigma')
probs = tf.placeholder(tf.float32, (None, 6, 12, 12, 2), name='probs')

dist = tfp.distributions.Mixture(cat=tfp.distributions.Categorical(probs=probs),
components=[tfp.distributions.Normal(loc=mu[:, :, :, :, 0], scale=sigma[:, :, :, :, 0]),
tfp.distributions.Normal(loc=mu[:, :, :, :, 1], scale=sigma[:, :, :, :, 1])])

with tf.Session() as sess:
batch_size = 4
m = np.random.normal(size=(batch_size, 6, 12, 12, 2))
s = np.random.normal(size=(batch_size, 6, 12, 12, 2))
s = np.exp(0.5*s)
p = np.random.uniform(0.0, 1.0, size=(batch_size, 6, 12, 12, 1))
p = np.concatenate([p, 1 - p], axis=-1)
feed_dict = {mu: m, sigma: s, probs: p}
sample = sess.run(dist.sample(), feed_dict=feed_dict)

Does someone can help ?

Thanks !!

Start by running the supplied TensorFlow session and Mixture example to reproduce the reported memory use. Trace Mixture.sample and its Normal components to determine where the excessive allocation occurs; done means identifying the cause and confirming a fix or documented limitation with a repeatable memory comparison.