code: def apply_mask(self, x, prev_output_tokens, original_batch_idxs):
      beam_size = x.shape[0] // original_batch_idxs.shape[0]
      original_batch_idxs = (
          original_batch_idxs.unsqueeze(-1)
          .repeat((1, beam_size))
          .flatten()
          .tolist()
      )
  
      mask = torch.full_like(x, -math.inf)
      for sent_i, (sent, batch_i) in enumerate(
          zip(prev_output_tokens, original_batch_idxs)
      ):
          mask[
              sent_i, :, self.prefix_allowed_tokens_fn(batch_i, sent)
          ] = 0
  
      return mask
  
imports: import math
  from typing import List, Optional
  import torch
  import torch.nn
  from fairseq.token_generation_constraints import ConstraintState, OrderedConstraintState, UnorderedConstraintState
  from torch import Tensor
label_desc: Use image node_id 2 for calling the PrefixConstrainedBeamSearch obj's underlying member method code with example usage: obj.apply_mask(x, prev_output_tokens, original_batch_idxs) and returns: mask