code: def prepare_masked_target_for_lengths(
      beam: Tensor,
      mask_idx: int,
      pad_idx: int,
      length_beam_size: int = 1,
  ) -> Tuple[Tensor, Tensor]:
      # beam : bsz X beam_size
      max_len = beam.max().item()
      bsz = beam.size(0)
      # length_mask[sample_length] will give a row vector of sample_length+1 ones
      # and rest zeros
      length_mask = torch.triu(
          torch.ones(
              max_len, max_len, device=beam.device, dtype=beam.dtype
          ),
          diagonal=1,
      )
      beam_indices = beam - 1
      length_mask = length_mask[beam_indices.reshape(-1)].reshape(
          bsz, length_beam_size, max_len
      )
      tgt_tokens = torch.zeros(
          bsz,
          length_beam_size,
          max_len,
          device=beam.device,
          dtype=beam.dtype,
      ).fill_(mask_idx)
      tgt_tokens = (
          1 - length_mask
      ) * tgt_tokens + length_mask * pad_idx
      tgt_tokens = tgt_tokens.view(bsz * length_beam_size, max_len)
      return tgt_tokens, length_mask
  
imports: from enum import Enum
  from typing import Dict, Optional, Tuple
  import torch
  from pytext.common.constants import TORCH_VERSION
  from pytext.config import ConfigBase
  from pytext.models.module import Module
  from pytext.torchscript.vocab import ScriptVocabulary
  from torch import Tensor
label_desc: Use image node_id 9 for calling a global function with example usage: prepare_masked_target_for_lengths(beam, mask_idx, pad_idx, length_beam_size) and returns: tgt_tokens,  length_mask