code: def compute_tdt_alphas_kernel(
      acts: torch.Tensor,
      duration_acts: torch.Tensor,
      denom: torch.Tensor,
      sigma: float,
      alphas: torch.Tensor,
      llForward: torch.Tensor,
      xlen: torch.Tensor,
      ylen: torch.Tensor,
      mlabels: torch.Tensor,  # [B]
      minibatch: int,
      maxT: int,
      maxU: int,
      alphabet_size: int,
      blank_: int,
      durations: torch.Tensor,
      num_durations: int,
  ):
      """
      Compute alpha (forward variable) probabilities over the transduction step.
  
      Args:
          acts: Tensor of shape [B, T, U, V] flattened. Represents the logprobs activation tensor for tokens.
          duration_acts: Tensor of shape [B, T, U, D] flattened. Represents the logprobs activation tensor for duration.
          denom: Tensor of shape [B, T, U] flattened. Represents the denominator of the logprobs activation tensor for tokens.
  
          alphas: Zero tensor of shape [B, T, U]. Will be updated inside the kernel with the forward variable
              probabilities.
          llForward: Zero tensor of shape [B]. Represents the log-likelihood of the forward pass.
              Returned as the forward pass loss that is reduced by the optimizer.
          xlen: Vector of length B which contains the actual acoustic sequence lengths in the padded
              activation tensor.
          ylen: Vector of length B which contains the actual target sequence lengths in the padded
              activation tensor.
          mlabels: Matrix of shape [B, U+1] (+1 here is due to <SOS> token - usually the RNNT blank).
              The matrix contains the padded target transcription that must be predicted.
          minibatch: Int representing the batch size.
          maxT: The maximum possible acoustic sequence length. Represents T in the logprobs tensor.
          maxU: The maximum possible target sequence length. Represents U in the logprobs tensor.
          alphabet_size: The vocabulary dimension V+1 (inclusive of RNNT blank).
          blank_: Index of the TDT blank token in the vocabulary. Must be the last token in the vocab.
  
      Updates:
          Kernel inplace updates the following inputs:
          -   alphas: forward variable scores.
          -   llForward: log-likelihood of forward variable.
      """
      # // launch B blocks, each block has U threads
      b = cuda.blockIdx.x  # // batch id
      u = cuda.threadIdx.x  # label id, u
      T = xlen[b]  # select AM length of current sample
      U = (
          ylen[b] + 1
      )  # select target length of current sample, +1 for the blank token
  
      labels: torch.Tensor = mlabels[
          b
      ]  # mb label start point, equivalent to mlabels + b * (maxU - 1)
      offset = b * maxT * maxU  # pointer indexing offset
  
      # alphas += offset # pointer offset, ignored since we explicitly add offset
  
      # Initilize alpha[b, t=0, u=0] for all b in B
      if u == 0:
          alphas[offset] = 0
  
      # sync until all alphas are initialized
      cuda.syncthreads()
  
      # Ordinary alpha calculations, broadcast across B=b and U=u
      # Look up forward variable calculation from rnnt_numpy.forward_pass()
      for n in range(1, T + U - 1):
          t = n - u
  
          if u == 0:
              # when u == 0, we only consider blank emissions.
              if t > 0 and t < T:
                  alphas[offset + t * maxU + u] = -INF
  
                  for i in range(
                      1, num_durations
                  ):  # skip 0 since blank emission has to advance by at least one
                      if t >= durations[i]:
                          alphas[
                              offset + t * maxU + u
                          ] = rnnt_helper.log_sum_exp(
                              alphas[
                                  offset + t * maxU + u
                              ],  # the current alpha value
                              alphas[
                                  offset + (t - durations[i]) * maxU + u
                              ]  # alpha(t - duration, u)
                              + logp(
                                  denom,
                                  acts,
                                  maxT,
                                  maxU,
                                  alphabet_size,
                                  b,
                                  t - durations[i],
                                  u,
                                  blank_,
                              )  # logp of blank emission
                              - sigma  #  logit under-normalization
                              + logp_duration(
                                  duration_acts,
                                  maxT,
                                  maxU,
                                  num_durations,
                                  b,
                                  t - durations[i],
                                  u,
                                  i,
                              ),  # logp of duration
                          )
                      else:
                          break  # since durations are in ascending order, when we encounter a duration that is too large, then
                          # there is no need to check larger durations after that.
  
          elif u < U:
              # when t == 0, we only consider the non-blank emission.
              if t == 0:
                  alphas[offset + u] = (
                      alphas[offset + u - 1]  # alpha(t, u - 1)
                      + logp(
                          denom,
                          acts,
                          maxT,
                          maxU,
                          alphabet_size,
                          b,
                          t,
                          u - 1,
                          labels[u - 1],
                      )  # logp of token emission
                      - sigma  # logit under-normalization
                      + logp_duration(
                          duration_acts,
                          maxT,
                          maxU,
                          num_durations,
                          b,
                          t,
                          u - 1,
                          0,
                      )  # t = 0, so it must be duration = 0. Therefore the last argument passed to logp_duration() is 0.
                  )
  
              # now we have t != 0 and u != 0, and we need to consider both non-blank and blank emissions.
              elif t > 0 and t < T:
                  no_emit = (
                      -INF
                  )  # no_emit stores the score for all blank emissions.
                  for i in range(1, num_durations):
                      if t >= durations[i]:
                          no_emit = rnnt_helper.log_sum_exp(
                              no_emit,  # current score
                              alphas[
                                  offset + (t - durations[i]) * maxU + u
                              ]  # alpha(t - duration, u)
                              + logp(
                                  denom,
                                  acts,
                                  maxT,
                                  maxU,
                                  alphabet_size,
                                  b,
                                  t - durations[i],
                                  u,
                                  blank_,
                              )  # logp of blank emission
                              - sigma  #  logit under-normalization
                              + logp_duration(
                                  duration_acts,
                                  maxT,
                                  maxU,
                                  num_durations,
                                  b,
                                  t - durations[i],
                                  u,
                                  i,
                              ),  # logp of duration
                          )
                      else:
                          break  # we can exit the loop early here, same as the case for u == 0 above.
  
                  emit = (
                      -INF
                  )  # emit stores the score for non-blank emissions.
                  for i in range(0, num_durations):
                      if t >= durations[i]:
                          emit = rnnt_helper.log_sum_exp(
                              emit,  # current score
                              alphas[
                                  offset
                                  + (t - durations[i]) * maxU
                                  + u
                                  - 1
                              ]  # alpha(t - duration, u - 1)
                              + logp(
                                  denom,
                                  acts,
                                  maxT,
                                  maxU,
                                  alphabet_size,
                                  b,
                                  t - durations[i],
                                  u - 1,
                                  labels[u - 1],
                              )  # logp of non-blank emission
                              - sigma  #  logit under-normalization
                              + logp_duration(
                                  duration_acts,
                                  maxT,
                                  maxU,
                                  num_durations,
                                  b,
                                  t - durations[i],
                                  u - 1,
                                  i,
                              ),  # logp of duration
                          )
                      else:
                          break  # we can exit the loop early here, same as the case for u == 0 above.
  
                  # combining blank and non-blank emissions.
                  alphas[
                      offset + t * maxU + u
                  ] = rnnt_helper.log_sum_exp(emit, no_emit)
  
          # sync across all B=b and U=u
          cuda.syncthreads()
  
      # After final sync, the forward log-likelihood can be computed as the summataion of
      # alpha(T - duration, U - 1) + logp(blank, duration | t - duration, U - 1), over different durations.
      if u == 0:
          # first we consider duration = 1
          loglike = (
              alphas[offset + (T - 1) * maxU + U - 1]
              + logp(
                  denom,
                  acts,
                  maxT,
                  maxU,
                  alphabet_size,
                  b,
                  T - 1,
                  U - 1,
                  blank_,
              )
              - sigma
              + logp_duration(
                  duration_acts,
                  maxT,
                  maxU,
                  num_durations,
                  b,
                  T - 1,
                  U - 1,
                  1,
              )
          )
  
          # then we add the scores for duration > 1, if such durations are possible given the audio lengths.
          for i in range(2, num_durations):
              if T >= durations[i]:
                  big_blank_loglike = (
                      alphas[offset + (T - durations[i]) * maxU + U - 1]
                      + logp(
                          denom,
                          acts,
                          maxT,
                          maxU,
                          alphabet_size,
                          b,
                          T - durations[i],
                          U - 1,
                          blank_,
                      )
                      - sigma
                      + logp_duration(
                          duration_acts,
                          maxT,
                          maxU,
                          num_durations,
                          b,
                          T - durations[i],
                          U - 1,
                          i,
                      )
                  )
                  loglike = rnnt_helper.log_sum_exp(
                      loglike, big_blank_loglike
                  )
              else:
                  break
  
          llForward[b] = loglike
  
file_path: NeMo/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt_kernel.py
imports: import math
  import torch
  from numba import cuda
  from nemo.collections.asr.parts.numba.rnnt_loss.utils import rnnt_helper
label_desc: Use image node_id 9 for calling a global function with example usage: compute_tdt_alphas_kernel(acts, duration_acts, denom, sigma, alphas, llForward, xlen, ylen, mlabels, minibatch, maxT, maxU, alphabet_size, blank_, durations, num_durations) without return types