def test_pad_mixture_dimensions_mixture(self):
gm = Mixture(cat=Categorical(probs=[[0.3, 0.7]]),
components=[
Normal(loc=[-1.0], scale=[1.0]),
Normal(loc=[1.0], scale=[0.5])
])
x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
x_pad = distribution_util.pad_mixture_dimensions(
x, gm, gm.cat, tensorshape_util.rank(gm.event_shape))
x_out, x_pad_out = self.evaluate([x, x_pad])
self.assertAllEqual(x_pad_out.shape, [2, 2])
self.assertAllEqual(x_out.reshape([-1]), x_pad_out.reshape([-1]))
def test_pad_mixture_dimensions_mixture(self):
with self.test_session() as sess:
gm = Mixture(cat=tf.distributions.Categorical(probs=[[0.3, 0.7]]),
components=[
tf.distributions.Normal(loc=[-1.0], scale=[1.0]),
tf.distributions.Normal(loc=[1.0], scale=[0.5])
])
x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
x_pad = distribution_util.pad_mixture_dimensions(
x, gm, gm.cat, gm.event_shape.ndims)
x_out, x_pad_out = sess.run([x, x_pad])
self.assertAllEqual(x_pad_out.shape, [2, 2])
self.assertAllEqual(x_out.reshape([-1]), x_pad_out.reshape([-1]))