docstr: Generate a pad mask tensor.
  
  # Args.
      input_lengths: (batch_size * beam_width,), input lengths
      memory_length: the length of key/value cache memory.
      init_step: int, initial step.
  # Return
      masked_tokens: BoolTensor,
          (batch_size * beam_width, memory_length),
          True if init_step + input_length[i] <= j <
              init_step + max_input_length,
          where i is a batch-beam index and j is a time step
          modulo by memory_length.
docstr_tok: ["Generate","a","pad","mask","tensor",".","#","Args",".","input_lengths",":","(","batch_size","*","beam_width",",",")",",","input","lengths","memory_length",":","the","length","of","key\/value","cache","memory",".","init_step",":","int",",","initial","step",".","#","Return","masked_tokens",":","BoolTensor",",","(","batch_size","*","beam_width",",","memory_length",")",",","True","if","init_step","+","input_length","[","i","]","<","=","j","<","init_step","+","max_input_length",",","where","i","is","a","batch-beam","index","and","j","is","a","time","step","modulo","by","memory_length","."]
code: def generate_pad_mask(
      self, input_lengths, memory_length, init_step=0
  ):
      """Generate a pad mask tensor.
  
      # Args.
          input_lengths: (batch_size * beam_width,), input lengths
          memory_length: the length of key/value cache memory.
          init_step: int, initial step.
      # Return
          masked_tokens: BoolTensor,
              (batch_size * beam_width, memory_length),
              True if init_step + input_length[i] <= j <
                  init_step + max_input_length,
              where i is a batch-beam index and j is a time step
              modulo by memory_length.
      """
      max_input_length = input_lengths.max()
      input_lengths = input_lengths.unsqueeze(1)
      shift = init_step % memory_length
      step_indices = torch.arange(
          init_step,
          init_step + memory_length,
          device=input_lengths.device,
      )
      step_indices = (
          step_indices.roll(shift)
          .unsqueeze(0)
          .tile(input_lengths.shape[0], 1)
      )
      masked_tokens = torch.logical_and(
          step_indices >= input_lengths,
          step_indices < init_step + max_input_length,
      )
      return masked_tokens
  
imports: import math
  from abc import abstractmethod
  from pathlib import Path
  from typing import List, Literal, Optional, Union
  import os
  import numpy
  from .None import comm
  from .None import profiler
  from .gpt import GptInitModelParameters
  from nebullvm.optional_modules.torch import torch
label_desc: Use image node_id 12 for calling the Gpt obj's underlying member method code with example usage: obj.generate_pad_mask(input_lengths, memory_length, init_step) and returns: masked_tokens