some confusions about soft-em

[yyht]

hi, since you waveglow propose to use a soft-em version of vqvae, the core implementation is:
"
def _square_distance(x, code_book):
x = tf.cast(x, tf.float32)
code_book = tf.cast(code_book, tf.float32)
x_sg = tf.stop_gradient(x)
x_norm_sq = tf.reduce_sum(tf.square(x_sg), axis=-1, keepdims=True) # [b, 1]
code_book_norm_sq = tf.reduce_sum(tf.square(code_book), axis=-1, keepdims=True) # [V, 1]
scalar_prod = tf.matmul(x_sg, code_book, transpose_b=True) # [b, V]
dist_sq = x_norm_sq + tf.transpose(code_book_norm_sq) - 2 * scalar_prod # [b, V]

return tf.cast(dist_sq, x.dtype.base_dtype)

dist_sq = _square_distance(x, code_book)
q = tf.stop_gradient(tf.nn.softmax(-.5 * dist_sq))
discrete = tf.one_hot(tf.argmax(-dist_sq, axis=-1), depth=bottleneck_size, dtype=code_book.dtype.base_dtype)
dense = tf.matmul(discrete, code_book)
dense = dense + x - tf.stop_gradient(x)
def _get_losses(x, x_mask, dense, dist_sq, q):
x = tf.cast(x, tf.float32)
x_mask = tf.cast(x_mask, tf.float32)
dense = tf.cast(dense, tf.float32)
dist_sq = tf.cast(dist_sq, tf.float32)
q = tf.cast(q, tf.float32)
disc_loss = tf.reduce_sum(tf.reduce_sum(tf.square(x - tf.stop_gradient(dense)), -1)*x_mask) / (1e-10+tf.reduce_sum(x_mask))
# # M-step
em_loss = -tf.reduce_sum(tf.reduce_sum(-.5 * dist_sq * q, -1)*x_mask) / (1e-10+tf.reduce_sum(x_mask))
return disc_loss, em_loss
disc_loss, em_loss = _get_losses(x, x_mask, dense, dist_sq, q)
"
, however, the tensor2tensor has a different implementation:
https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/layers/vq_discrete.py

1. multisample to get mean of soft-alignment
2. when calculate em-loss, it has a different loss funtion type compare to your "M-step"
   .
   Could you hepl me with it?