trainer_base.py

import itertools
from dataclasses import dataclass
import random

import hydra.utils
import lightning as L
import numpy as np
import torch
import torch.nn.functional as F
import transformers
# import mup

import dataloader
import metrics
import models
import utils

from samplers import _get_sampler


@dataclass
class Loss:
  loss: torch.FloatTensor
  nlls: torch.FloatTensor
  reconstruction_loss: torch.FloatTensor
  num_tokens: torch.FloatTensor


class LogLinear(torch.nn.Module):
  def __init__(self, alpha_0=1):
    super().__init__()
    self.eps = 1e-3  # To be consistent with SEDD: https://github.com/louaaron/Score-Entropy-Discrete-Diffusion/blob/0605786da5ccb5747545e26d66fdf477187598b6/noise_lib.py#L56
    self.alpha_0 = alpha_0

  def forward(self, t):
    t = (1 - self.eps) * t
    alpha_t = self.alpha_0 * (1 - t)
    dalpha_t = - self.alpha_0 * (1 - self.eps)
    return dalpha_t, alpha_t


def sample_categorical(categorical_probs):
  gumbel_norm = (
    1e-10
    - (torch.rand_like(categorical_probs) + 1e-10).log())
  return (categorical_probs / gumbel_norm).argmax(dim=-1)


def _unsqueeze(x, reference):
  return x.view(
    * x.shape,
    * ((1,) * (len(reference.shape) - len(x.shape))))


class TrainerBase(L.LightningModule):
  def __init__(self, config, tokenizer):
    super().__init__()
    self.save_hyperparameters()
    self.config = config
    if hasattr(self.config.algo, 'loss_type'):
      self.loss_type = config.algo.loss_type
    self.tokenizer = tokenizer
    self.vocab_size = len(tokenizer)
    if (not hasattr(tokenizer, 'mask_token')
        or tokenizer.mask_token is None):
      self.mask_index = self.vocab_size
      self.vocab_size += 1
    else:
      self.mask_index = tokenizer.mask_token_id
    self.sampler = self.config.sampling.predictor
    self.antithetic_sampling = self.config.training.antithetic_sampling
    self.parameterization = self.config.algo.parameterization
    if self.config.algo.backbone in ['dit', 'dit_legacy']:
      self.backbone = models.dit.DiT(
        self.config, vocab_size=self.vocab_size)
    elif self.config.algo.backbone in ['esolm_dit', 'esolm_dit_legacy']:
      self.backbone = models.dit.EsoLMDiT(
        self.config, vocab_size=self.vocab_size,
        mask_index=self.mask_index)  # mask_index is defined in the child class
    elif self.config.algo.backbone == 'hf_dit':
      self.backbone = transformers.AutoModelForMaskedLM.from_pretrained(
        config.eval.checkpoint_path, trust_remote_code=True)
    elif self.config.algo.backbone == 'vdlm_dit':
      self.backbone = models.dit.VDLMDiT(
        self.config, vocab_size=self.vocab_size, 
        mask_index=self.mask_index)
    elif self.config.algo.backbone == 'vdlm_dit_vq':
      self.backbone = models.dit.VDLMDiTVQ(
        self.config, vocab_size=self.vocab_size, 
        mask_index=self.mask_index)
    elif self.config.algo.backbone == 'llada_dit':
      self.backbone = transformers.AutoModel.from_pretrained(
        self.config.algo.model_path, trust_remote_code=True, torch_dtype=torch.bfloat16)
    elif self.config.algo.backbone == 'mlp':
      # self.backbone = models.mlp.MLPDenoiser(
      self.backbone = models.mlp.MLP(
        self.config, vocab_size=self.vocab_size)
    elif self.config.algo.backbone == 'vdlm_mlp':
      self.backbone = models.mlp.VDLMMLP(
        self.config, vocab_size=self.vocab_size)

    self.T = self.config.algo.T
    self.num_tokens = self.config.model.length
    self.softplus = torch.nn.Softplus()
    self.noise = LogLinear()
    self.p_nucleus = self.config.sampling.p_nucleus
    self.metrics = metrics.Metrics(
      gen_ppl_eval_model_name_or_path=self.config.eval.gen_ppl_eval_model_name_or_path,
      eval_ppl_batch_size=self.config.eval.perplexity_batch_size)

    if self.config.training.ema > 0:
      self.ema = models.ema.ExponentialMovingAverage(
        self._get_parameters(),
        decay=self.config.training.ema)
    else:
      self.ema = None

    self.lr = self.config.optim.lr
    self.sampling_eps = float(self.config.training.sampling_eps)
    self.time_conditioning = self.config.algo.time_conditioning
    self.neg_infinity = -1000000.0
    self.fast_forward_epochs = None
    self.fast_forward_batches = None
    self.train_start_file_idx = 0

  def setup(self, stage=None):
    # different randomness for different ranks
    # does not affect dataloading seed
    del stage
    new_seed = self.config.seed + self.trainer.global_rank
    torch.manual_seed(new_seed)
    np.random.seed(new_seed)
    random.seed(new_seed)

  def _validate_configuration(self):
    assert self.config.algo.backbone in {'dit', 'dit_legacy', 'hf_dit', 'llada_dit',
                                         'esolm_dit', 'esolm_dit_legacy',
                                         'vdlm_dit', 'vdlm_dit_vq', 
                                         'mlp', 'vdlm_mlp'}
    if self.config.algo.parameterization == 'ar':
      assert not self.config.algo.time_conditioning
      assert self.config.prior.type == 'none'

    if self.parameterization in {'score', 'mean'}:
      assert self.time_conditioning
    if self.T > 0:
      assert self.parameterization != 'score'

  def to(self, *args, **kwargs):
    self = super().to(*args, **kwargs)
    self.metrics.to(*args, **kwargs)
    return self

  def q_xt(self, x, alpha_t):
    raise NotImplementedError

  def _get_parameters(self):
    return itertools.chain(self.backbone.parameters(),
                           self.noise.parameters())

  def _eval_mode(self):
    if self.ema:
      self.ema.store(self._get_parameters())
      self.ema.copy_to(self._get_parameters())
    self.backbone.eval()
    self.noise.eval()

  def _train_mode(self):
    if self.ema:
      self.ema.restore(self._get_parameters())
    self.backbone.train()
    self.noise.train()

  def on_load_checkpoint(self, checkpoint):
    if self.ema and 'ema' in checkpoint.keys():
      self.ema.load_state_dict(checkpoint['ema'])
    # Copied from:
    # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/datamodules/language_modeling_hf.py#L41
    self.fast_forward_epochs = checkpoint['loops'][
      'fit_loop']['epoch_progress']['current']['completed']
    self.fast_forward_batches = checkpoint['loops'][
      'fit_loop']['epoch_loop.batch_progress'][
        'current']['completed']
    if 'train_start_file_idx' in checkpoint.keys():  # Backward compatibility
      self.train_start_file_idx = checkpoint['train_start_file_idx']

  def on_save_checkpoint(self, checkpoint):
    if self.ema:
      checkpoint['ema'] = self.ema.state_dict()
    checkpoint['train_start_file_idx'] = self.train_start_file_idx
    # Copied from:
    # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/tasks/seq.py
    # ['epoch_loop.batch_progress']['total']['completed']
    # is 1 iteration behind, so we're using the optimizer's progress.
    checkpoint['loops']['fit_loop'][
      'epoch_loop.batch_progress']['total'][
        'completed'] = checkpoint['loops']['fit_loop'][
          'epoch_loop.automatic_optimization.optim_progress'][
            'optimizer']['step']['total'][
              'completed'] * self.trainer.accumulate_grad_batches
    checkpoint['loops']['fit_loop'][
      'epoch_loop.batch_progress']['current'][
        'completed'] = checkpoint['loops']['fit_loop'][
          'epoch_loop.automatic_optimization.optim_progress'][
            'optimizer']['step']['current'][
              'completed'] * self.trainer.accumulate_grad_batches
    # _batches_that_stepped tracks the number of global steps,
    # not the number of local steps, so we don't multiply with
    # self.trainer.accumulate_grad_batches here.
    checkpoint['loops']['fit_loop'][
      'epoch_loop.state_dict'][
        '_batches_that_stepped'] = checkpoint['loops']['fit_loop'][
          'epoch_loop.automatic_optimization.optim_progress'][
            'optimizer']['step']['total']['completed']
    if 'sampler' not in checkpoint.keys():
      checkpoint['sampler'] = {}
    if hasattr(self.trainer.train_dataloader.sampler,
               'state_dict'):
      sampler_state_dict = self.trainer.\
        train_dataloader.sampler.state_dict()
      checkpoint['sampler'][
        'random_state'] = sampler_state_dict.get(
          'random_state', None)
    else:
      checkpoint['sampler']['random_state'] = None

  def on_train_start(self):
    if self.ema:
      self.ema.move_shadow_params_to_device(self.device)
    # Adapted from:
    # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/datamodules/language_modeling_hf.py
    updated_dls = []
    for dl in self.trainer.fit_loop._combined_loader.flattened:
      updated_dls.append(
        torch.utils.data.DataLoader(
          dl.dataset,
          batch_size=self.config.loader.batch_size,
          num_workers=self.config.loader.num_workers,
          pin_memory=self.config.loader.pin_memory,
          # sampler=dl_sampler,
          shuffle=False,
          persistent_workers=True))
    self.trainer.fit_loop._combined_loader.flattened = updated_dls

  def optimizer_step(self, *args, **kwargs):
    super().optimizer_step(*args, **kwargs)
    if self.ema:
      self.ema.update(self._get_parameters())

  def _process_sigma(self, sigma):
    raise NotImplementedError

  def _process_model_output(self, model_output, xt, sigma):
    raise NotImplementedError

  @torch.no_grad()
  def augment_batch_cfg(self, batch, prompt_index, sigma=None, sort_idx=None):
    if self.config.sampling.cfg > 0.:
      if isinstance(prompt_index, list):
        prompt_index = torch.tensor(prompt_index)
      if prompt_index.ndim < 2:
        prompt_index = prompt_index.unsqueeze(
          0).repeat(batch.shape[0], 1)
      assert prompt_index.shape[1] == batch.shape[1], f"Expected prompt_index length {prompt_index.shape[1]} to be equal to sequence length {batch.shape[1]}"
      assert prompt_index.shape[0] == batch.shape[0], f"Expected prompt_index batch size {prompt_index.shape[0]} to be equal to sequence batch size {batch.shape[0]}"
      un_batch = batch.clone()
      prior_batch = self.prior_sample(*batch.shape)
      un_batch[prompt_index] = prior_batch[prompt_index]
      batch = torch.cat([batch, un_batch])
      if sigma is not None:
        sigma = torch.cat([sigma, sigma])
      if sort_idx is not None:
        sort_idx = torch.cat([sort_idx, sort_idx])
    return batch, sigma, sort_idx

  @torch.no_grad()
  def reduce_batch_cfg(self, batch, logits):
    if self.config.sampling.cfg > 0.:
      logits, un_logits = torch.chunk(logits, 2, dim=0)
      logits = un_logits + \
          (self.config.sampling.cfg + 1) * \
          (logits - un_logits)
    return logits

  def forward(self, xt, sigma, sort_idx=None, x0=None, latent=None, attn_mask=None, prompt_index=None, dynamic=False):
    sigma = self._process_sigma(sigma)
    xt, sigma, sort_idx = self.augment_batch_cfg(xt, prompt_index, sigma=sigma, sort_idx=sort_idx)
    with torch.amp.autocast('cuda', dtype=torch.float32):
      logits = self.backbone(xt, sigma, sort_idx, x0, latent=latent, attn_mask=attn_mask, dynamic=dynamic)
    logits = self.reduce_batch_cfg(xt, logits)
    return self._process_model_output(
      model_output=logits, xt=xt, sigma=sigma)

  def on_train_epoch_start(self):
    self.metrics.reset()
    self.metrics.to(self.device)
    assert self.metrics.train_nlls.nll.mean_value == 0
    assert self.metrics.train_nlls.nll.weight == 0

  def training_step(self, batch, batch_idx):
    torch.compiler.cudagraph_mark_step_begin()
    current_accumulation_step = (
      batch_idx % self.trainer.accumulate_grad_batches)
    input_tokens = batch['input_ids']
    if 'file_idx' in batch.keys():
      self.train_start_file_idx = batch['file_idx'].max(
      ).cpu().item()

    if (self.config.algo.name != 'ar'
        and torch.rand(1) < self.config.training.ssl_ratio):
      length = torch.randint(1, input_tokens.shape[-1] + 1, (1,))
      input_tokens = input_tokens[:, :length]
    attention_mask = torch.ones_like(input_tokens)

    losses = self._loss(input_tokens, attention_mask,
                        current_accumulation_step,
                        train_mode=True)
    if torch.isnan(losses.loss).any():
      raise ValueError('Loss is nan')
    self.metrics.update_train(losses.nlls,
                              losses.reconstruction_loss,
                              losses.num_tokens)

    for key in losses.__dict__.keys():
      if isinstance(getattr(losses, key), torch.Tensor):
        self.log(name=f'trainer/{key}',
                value=getattr(losses, key).item(),
                on_step=True,
                on_epoch=False,
                sync_dist=True)
      elif isinstance(getattr(losses, key), dict):
        for k, v in getattr(losses, key).items():
          self.log(name=f'trainer/{key}/{k}',
                  value=v.item(),
                  on_step=True,
                  on_epoch=False,
                  sync_dist=True)
    self.log(name='trainer/train_start_file_idx',
            value=self.train_start_file_idx,
            on_step=True,
            on_epoch=False,
            sync_dist=True)
    return losses.loss

  def on_train_epoch_end(self):
    for k, v in self.metrics.train_nlls.items():
      self.log(name=k, value=v.compute(), on_step=False,
               on_epoch=True, sync_dist=True)

  def on_validation_epoch_start(self):
    self.metrics.reset()
    self._eval_mode()
    assert self.metrics.valid_nlls.nll.mean_value == 0
    assert self.metrics.valid_nlls.nll.weight == 0

  def validation_step(self, batch, batch_idx):
    del batch_idx
    input_tokens = batch['input_ids']
    attention_mask = torch.ones_like(input_tokens)
    losses = self._loss(input_tokens, attention_mask)
    self.metrics.update_valid(losses.nlls,
                              losses.reconstruction_loss,
                              losses.num_tokens)
    return losses.loss

  def on_validation_epoch_end(self):
    for k, v in self.metrics.valid_nlls.items():
      self.log(name=k, value=v.compute(), on_step=False,
               on_epoch=True, sync_dist=True)
    if ((self.config.eval.compute_perplexity_on_sanity
         or not self.trainer.sanity_checking)
         and self.config.eval.generate_samples):
      samples, text_samples = None, None
      num_sample_batches = self.config.sampling.num_samples // (
          self.trainer.num_nodes  * self.trainer.num_devices
          * self.config.loader.eval_batch_size)
      for _ in range(max(num_sample_batches, 1)):
        samples = self.generate_samples(
          num_samples=self.config.loader.eval_batch_size)
        self.metrics.record_entropy(samples)  
        # Decode the samples to be re-tokenized by eval model
        text_samples = self.tokenizer.batch_decode(samples)
        if self.config.eval.compute_generative_perplexity:
          self.metrics.record_generative_perplexity(
            text_samples, self.num_tokens, self.device)
      if text_samples is not None:
        if self.trainer.global_rank == 0 and hasattr(
          self.trainer.logger, 'log_table'):
          # Log the last generated samples
          text_samples = text_samples[
            : self.config.sampling.num_log_samples]
          self.trainer.logger.log_table(
            key=f'samples@global_step{self.global_step}',
            columns=['Generated Samples'],
            data=[[s] for s in text_samples])
        if self.config.eval.compute_generative_perplexity:
          self.log('val/gen_ppl',
                   self.metrics.gen_ppl.compute(),
                   on_epoch=True,
                   on_step=False,
                   sync_dist=True)
          self.log('val/sample_entropy',
                   self.metrics.sample_entropy.compute(),
                   on_epoch=True,
                   on_step=False,
                   sync_dist=True)
    self._train_mode()

  def configure_optimizers(self):
    optimizer = torch.optim.AdamW(
      self._get_parameters(),
      lr=self.config.optim.lr,
      betas=(self.config.optim.beta1,
             self.config.optim.beta2),
      eps=self.config.optim.eps,
      weight_decay=self.config.optim.weight_decay)

    scheduler = hydra.utils.instantiate(
      self.config.lr_scheduler, optimizer=optimizer)
    scheduler_dict = {'scheduler': scheduler,
                      'interval': 'step',
                      'monitor': 'val/loss',
                      'name': 'trainer/lr'}
    return [optimizer], [scheduler_dict]

  @torch.no_grad()
  def generate_samples(self, num_samples, eps=1e-5, condition=None):
    """Generate samples from the model."""
    # Lightning auto-casting is not working in this method for some reason
    sampler = _get_sampler(
      self.config, self, self.tokenizer)
    return sampler(
      num_samples=num_samples,
      eps=eps,
      condition=condition)

  def restore_model_and_sample(self, eps=1e-5, condition=None):
    """Generate samples from the model."""
    # Lightning auto-casting is not working in this method for some reason
    self._eval_mode()
    samples = self.generate_samples(
      num_samples=self.config.loader.eval_batch_size,
      eps=eps,
      condition=condition)
    self._train_mode()
    return samples

  def _process_model_input(self, x0, valid_tokens):
    raise NotImplementedError

  def nll(self, input_tokens, output_tokens,
          current_accumulation_step=None, train_mode=False):
    raise NotImplementedError

  def _loss(self, x0, valid_tokens,
            current_accumulation_step=None,
            train_mode=False):
    (input_tokens, output_tokens,
     valid_tokens) = self._process_model_input(
       x0, valid_tokens)
    loss = self.nll(input_tokens, output_tokens,
                    current_accumulation_step, train_mode)
    assert loss.ndim == 2

    nlls = (loss * valid_tokens).sum()
    num_tokens = valid_tokens.sum()
    token_nll = nlls / num_tokens

    return Loss(loss=token_nll,
                nlls=nlls,
                reconstruction_loss=torch.tensor(0),
                num_tokens=num_tokens)

  ### TODO: test that and implement with resume etc.
  # def predict_dataloader(self):
  #   world_size = (self.trainer.num_nodes
  #                 * self.trainer.num_devices)
  #   total_batches = math.ceil(
  #     self.config.sampling.num_samples / (
  #       world_size * self.config.loader.eval_batch_size))
  #   # Ensure the dataset length is a multiple of world size 
  #   # so each rank receives at least one sample with the
  #   # unrepeated distributed sampler.
  #   return torch.utils.data.DataLoader(
  #     torch.utils.data.TensorDataset(
  #       torch.arange(total_batches * world_size)),
  #     batch_size=1,
  #     num_workers=0,
  #     pin_memory=False,
  #     shuffle=False)

  # def on_predict_start(self):
  #   self.metrics.gen_ppl.reset()
  #   self.metrics.sample_entropy.reset()
  #   self._eval_mode()

  # def predict_step(self, batch, batch_idx, dataloader_idx=0):
  #   del batch, batch_idx, dataloader_idx
  #   samples = self.generate_samples(
  #     num_samples=self.config.loader.eval_batch_size)
    
  #   self.metrics.record_entropy(samples)
  #   # Decode the samples to be re-tokenized by eval model
  #   text_samples = self.tokenizer.batch_decode(samples)
  #   if self.config.eval.compute_generative_perplexity:
  #     self.metrics.record_generative_perplexity(
  #       text_samples, self.num_tokens, self.device)
  #   return text_samples


class Diffusion(TrainerBase):
  def _validate_configuration(self):
    super()._validate_configuration()
    assert self.config.sampling.noise_removal in {
      'none', 'ancestral', 'greedy'}
    assert self.loss_type in {'elbo', 'low_var', 'delbo'}
    if self.config.sampling.noise_removal == 'greedy':
      assert self.sampler != 'analytic'
      assert self.parameterization in {'mean', 'subs'}

  def _process_model_input(self, x0, valid_tokens):
    return x0, None, valid_tokens

  def _process_sigma(self, sigma):
    assert sigma.ndim == 2
    sigma = sigma.mean(-1).squeeze()
    if sigma.ndim == 0:
      sigma = sigma.unsqueeze(0)
    if not self.time_conditioning:
      sigma = torch.zeros_like(sigma)
    assert sigma.ndim == 1, sigma.shape
    return sigma

  def _sample_t(self, n, accum_step):
    if accum_step is not None:
      # During training
      batch_dim = n
      n = self.config.loader.global_batch_size
    _eps_t = torch.rand(n, device=self.device)
    if self.antithetic_sampling:
      offset = torch.arange(n, device=self.device) / n
      _eps_t = (_eps_t / n + offset) % 1
    t = (1 - self.sampling_eps) * _eps_t + self.sampling_eps
    if accum_step is not None:
      t = t.chunk(self.trainer.num_nodes)[
          self.trainer.node_rank]
      t = t.chunk(self.trainer.num_devices)[
          self.trainer.local_rank]
      t = t.chunk(self.trainer.accumulate_grad_batches)[
        accum_step]
      # corner case for the last datapoint
      t = t[:batch_dim]
    return t

  def _sigma_from_alphat(self, alpha_t):
    return -torch.log(alpha_t)

  def _reconstruction_loss(self, x0):
    t0 = torch.zeros(1, x0.shape[0], dtype=self.dtype,
                     device=self.device)
    sigma_t0 = self._sigma_from_alphat(self.noise(t0)[1])
    model_output_t0 = self.forward(x0, sigma_t0)
    return - torch.gather(input=model_output_t0,
                          dim=-1,
                          index=x0[:, :, None]).squeeze(-1)

  def nll_per_token(self, model_output, xt, x0, alpha_t,
                    dalpha_t, low_var):
    raise NotImplementedError

  def nll(self, x0, output_tokens,
          current_accumulation_step=None, train_mode=False):
    del output_tokens
    t = self._sample_t(x0.shape[0],
                       current_accumulation_step)
    assert t.shape[0] == x0.shape[0]
    if self.T > 0:
      t = (t * self.T).to(torch.int)
      t = t / self.T
      # t \in {1/T, 2/T, ..., 1}
      t += (1 / self.T)

    dalpha_t, alpha_t = self.noise(t)
    alpha_t = alpha_t.unsqueeze(-1)
    assert alpha_t.ndim == 2
    sigma = self._sigma_from_alphat(alpha_t)

    xt = self.q_xt(x0, alpha_t)
    # print("xt.shape", xt.shape)
    # print("Masked xt", (xt == self.mask_index).sum() / xt.numel())
    log_x_theta = self.forward(xt, sigma=sigma)
    # print("log_x_theta.shape", log_x_theta.shape)
    utils.print_nans(log_x_theta, 'model_output')
    return self.nll_per_token(
      log_x_theta=log_x_theta,
      xt=xt,
      x0=x0,
      alpha_t=alpha_t,
      dalpha_t=dalpha_t,
      low_var=train_mode and self.loss_type == 'low_var')

  def _get_score(self, **kwargs):
    del kwargs
    raise NotImplementedError

  def _denoiser_update(self, x, t):
    raise NotImplementedError

  def _analytic_update(self, x, t, dt):
    raise NotImplementedError

  def _ancestral_update(self, x, t, dt, p_x0, noise_removal_step):
    raise NotImplementedError


class AbsorbingState(Diffusion):
  def __init__(self, config, tokenizer):
    self.subs_masking = config.algo.subs_masking
    super().__init__(config, tokenizer)
    self.save_hyperparameters()

  def _validate_configuration(self):
    super()._validate_configuration()
    if self.parameterization in {'score', 'mean'}:
      assert self.time_conditioning
    assert not (self.parameterization == 'mean'
                and self.T == 0)
    if self.T > 0:
      assert self.parameterization in {'mean', 'subs'}
    if self.subs_masking:
      assert self.parameterization == 'mean'

  def q_xt(self, x, alpha_t):
    """Computes the noisy sample xt.

    Args:
      x: int torch.Tensor with shape (batch_size,
          diffusion_model_input_length), input. 
      alpha_t: float torch.Tensor with shape (batch_size, 1).
    """
    move_indices = torch.rand(
      * x.shape, device=x.device) < 1 - alpha_t
    xt = torch.where(move_indices, self.mask_index, x)
    return xt

  def prior_sample(self, *batch_dims):
    return self.mask_index * torch.ones(
      * batch_dims, dtype=torch.int64, device=self.device)


class UniformState(Diffusion):
  def _validate_configuration(self):
    super()._validate_configuration()
    assert self.time_conditioning
    assert self.parameterization == 'mean'
    if self.config.algo.name != 'distillation':
      assert self.T == 0

  def q_xt(self, x, alpha_t):
    """Computes the noisy sample xt.

    Args:
      x: int torch.Tensor with shape (batch_size,
          diffusion_model_input_length), input.
      move_chance: float torch.Tensor with shape
        (batch_size, 1).
    """
    move_indices = torch.rand(
      *x.shape, device=x.device) < 1 - alpha_t
    uniform_tensor = torch.randint(
      0, self.vocab_size, x.shape, device=x.device)
    xt = torch.where(move_indices, uniform_tensor, x)
    return xt

  def prior_sample(self, *batch_dims):
    return torch.randint(
      0, self.vocab_size, batch_dims, dtype=torch.int64,
      device=self.device)