import tensorflow as tf

          tf.random.categorical(
              logits[None], num_samples=self._sample_batch_size))
      mask = tf.cast(mask, tf.float32)[:, None]
      relabelled_tasks = mask * orig_tasks + (1 - mask) * relabelled_tasks
