import tensorflow as tf

          dist /= tf.reshape(self.hparams.scales,
                             [1, 1, self.hparams.moe_num_experts])
        nearest_idx = tf.argmax(-dist, axis=-1)
      nearest_hot = tf.one_hot(nearest_idx, self.hparams.block_v_size)
