import tensorflow as tf

          moving_averages.assign_moving_average(
              self.ema_count,
              tf.reduce_sum(
                  tf.reshape(
                      x_means_hot,
                      shape=[-1, self.hparams.num_blocks,
                             self.hparams.block_v_size]),
                  axis=0),
              self.hparams.decay,
              zero_debias=False)

      dw = tf.matmul(
          tf.transpose(x_means_hot, perm=[1, 2, 0]),
          tf.transpose(x_reshaped, perm=[1, 0, 2]))

      updated_ema_means = \
          moving_averages.assign_moving_average(
              self.ema_means, dw, self.hparams.decay,
              zero_debias=False)
      n = tf.reduce_sum(updated_ema_count, axis=-1, keep_dims=True)
      updated_ema_count = ((updated_ema_count + self.hparams.epsilon) / (
          n + 2**self.hparams.z_size * self.hparams.epsilon) * n)
      updated_ema_means = updated_ema_means / tf.expand_dims(
          updated_ema_count, axis=-1)

      with tf.control_dependencies([e_loss]):
