def testVQDiscreteUnbottlenck(self):
     x = tf.constant([[1, 0, 0, 0], [0, 0, 1, 0]], dtype=tf.int32)
     x_means = discretization.vq_discrete_unbottleneck(x, hidden_size=3)
     with self.test_session() as sess:
         tf.global_variables_initializer().run()
         x_means_eval = sess.run(x_means)
         self.assertEqual(np.shape(x_means_eval), (2, 3))
 def testVQDiscreteUnbottlenck(self):
   x = tf.constant([[1, 0, 0, 0], [0, 0, 1, 0]], dtype=tf.int32)
   x_means = discretization.vq_discrete_unbottleneck(x, hidden_size=3)
   with self.test_session() as sess:
     tf.global_variables_initializer().run()
     x_means_eval = sess.run(x_means)
     self.assertEqual(np.shape(x_means_eval), (2, 3))
 def testVQDiscreteUnbottlenck(self):
     x = tf.constant([[1, 0, 0, 0], [0, 0, 1, 0]], dtype=tf.int32)
     x_means = discretization.vq_discrete_unbottleneck(x, model_d=3)
     self.evaluate(tf.global_variables_initializer())
     x_means_eval = self.evaluate(x_means)
     self.assertEqual(np.shape(x_means_eval), (2, 3))