code: def spec_augment_kernel(
      x: torch.Tensor,
      x_len: torch.Tensor,
      freq_starts: torch.Tensor,
      freq_widths: torch.Tensor,
      time_starts: torch.Tensor,
      time_widths: torch.Tensor,
      mask_value: float,
  ):
      """
      Numba CUDA kernel to perform SpecAugment in-place on the GPU.
      Parallelize over freq and time axis, parallel threads over batch.
      Sequential over masks (adaptive in time).
  
      Args:
          x: Pytorch tensor of shape [B, F, T] with the acoustic features.
          x_len: Pytorch tensor of shape [B] with the lengths of the padded sequence.
          freq_starts: Pytorch tensor of shape [B, M_f] with the start indices of freq masks.
          freq_widths: Pytorch tensor of shape [B, M_f] with the width of freq masks.
          time_starts: Pytorch tensor of shape [B, M_t] with the start indices of time masks.
          time_widths: Pytorch tensor of shape [B, M_t] with the width of time masks.
          mask_value: Float value that will be used as mask value.
      """
      f = cuda.blockIdx.x  # indexes the Freq dim
      t = cuda.blockIdx.y  # indexes the Time dim
      tid = cuda.threadIdx.x  # index of the current mask
      threads_per_block = cuda.blockDim.x
  
      # Compute the number of masks over freq axis
      len_f = freq_starts.shape[1]
      # For all samples in the batch, apply the freq mask
      for bidx in range(0, x.shape[0], threads_per_block):
          # Resolve the index of the batch (case where more masks than MAX_THREAD_BUFFER)
          bm_idx = bidx + tid
  
          # Access mask only if valid sample id in batch
          if bm_idx < x.shape[0]:
              # For `len_f` number of freq masks that must be applied
              for fidx in range(0, len_f):
                  # Access the start index and width of this freq mask
                  f_start = freq_starts[bm_idx, fidx]
                  f_width = freq_widths[bm_idx, fidx]
  
                  # If block idx `f` >= start and < (start + width) of this freq mask
                  if f >= f_start and f < (f_start + f_width):
                      x[bm_idx, f, t] = mask_value
  
      # Compute the number of masks over time axis
      len_t = time_starts.shape[1]
      # For all samples in the batch, apply the time mask
      for b_idx in range(0, x.shape[0], threads_per_block):
          # Resolve the index of the batch (case where more masks than MAX_THREAD_BUFFER)
          bm_idx = b_idx + tid
  
          # Access mask only if valid sample id in batch
          if bm_idx < x.shape[0]:
              # For `len_t` number of freq masks that must be applied
              for tidx in range(0, len_t):
                  # Access the start index and width of this time mask
                  t_start = time_starts[bm_idx, tidx]
                  t_width = time_widths[bm_idx, tidx]
  
                  # If block idx `t` >= start and < (start + width) of this time mask
                  if t >= t_start and t < (t_start + t_width):
                      # Current block idx `t` < current seq length x_len[b]
                      # This ensure that we mask only upto the length of that sample
                      # Everything after that index is padded value so unnecessary to mask
                      if t < x_len[bm_idx]:
                          x[bm_idx, f, t] = mask_value
  
file_path: NeMo/nemo/collections/asr/parts/numba/spec_augment/spec_aug_numba.py
imports: import torch
  import torch.nn
  from numba import cuda
  from nemo.core.classes import Typing, typecheck
  from nemo.core.neural_types import LengthsType, NeuralType, SpectrogramType
  from nemo.utils import logging
label_desc: Use image node_id 1 for calling a global function with example usage: spec_augment_kernel(x, x_len, freq_starts, freq_widths, time_starts, time_widths, mask_value) without return types