import tensorflow as tf

    """Slice encoder hidden state into block_dim.

    Args:
        x: Encoder hidden state of shape [-1, hidden_size].

    Returns:
        Sliced states of shape [-1, num_blocks, block_dim].
    """
    x_sliced = tf.reshape(
        x, shape=[-1, self.hparams.num_blocks, self.hparams.block_dim])
    return x_sliced

  def nearest_neighbor(self, x, means):
    """Find the nearest element in means to elements in x.

    Args:
