code: def compute_multiblank_betas_kernel(
      acts: torch.Tensor,
      denom: torch.Tensor,
      sigma: float,
      betas: torch.Tensor,
      llBackward: torch.Tensor,
      xlen: torch.Tensor,
      ylen: torch.Tensor,
      mlabels: torch.Tensor,  # [B, U]
      minibatch: int,
      maxT: int,
      maxU: int,
      alphabet_size: int,
      blank_: int,
      big_blank_duration: torch.Tensor,
      num_big_blanks: int,
  ):
      """
      Compute beta (backward variable) probabilities for multi-blank transducer loss
          (https://arxiv.org/pdf/2211.03541).
  
      Args:
          acts: Tensor of shape [B, T, U, V + 1 + num-big-blanks] flattened.
              Represents the logprobs activation tensor.
          denom: Tensor of shape [B, T, U] flattened. Represents the denominator
              of the logprobs activation tensor across entire vocabulary.
          sigma: Hyper-parameter for logit-undernormalization technique for
              training multi-blank transducers.
          betas: Zero tensor of shape [B, T, U]. Will be updated inside the kernel
              with the backward variable probabilities.
          llBackward: Zero tensor of shape [B]. Represents the log-likelihood
              of the backward pass. Returned as the backward 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 RNNT standard blank token in the vocabulary.
          big_blank_durations: Vector of supported big blank durations of the model.
          num_big_blanks: Number of big blanks of the model.
  
      Updates:
          Kernel inplace updates the following inputs:
          -   betas: backward variable scores.
          -   llBackward: log-likelihood of backward 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
  
      # Note: just like the alphas, because of the logit under-normalization, everytime
      # logp() is called, it is always followed by a `-sigma` term.
  
      # Initilize beta[b, t=T-1, u=U-1] for all b in B with
      # log_probs[b, t=T-1, u=U-1, blank]
      if u == 0:
          betas[offset + (T - 1) * maxU + U - 1] = (
              logp(
                  denom,
                  acts,
                  maxT,
                  maxU,
                  alphabet_size,
                  b,
                  T - 1,
                  U - 1,
                  blank_,
              )
              - sigma
          )
  
      # sync until all betas are initialized
      cuda.syncthreads()
  
      # Ordinary beta calculations, broadcast across B=b and U=u
      # Look up backward variable calculation from rnnt_numpy.backward_pass()
      for n in range(T + U - 2, -1, -1):
          t = n - u
  
          if u == (U - 1):
              # for t in reversed(range(T - 1)) step to initialize betas[b, t, U-1]
              if t >= 0 and t < (T - 1):
                  # beta[t, U - 1] = beta[t + 1, U - 1] * p(blank | t, U - 1) / exp(sigma)
                  # this part is the same as regular RNN-T.
                  betas[offset + t * maxU + U - 1] = (
                      betas[offset + (t + 1) * maxU + U - 1]
                      + logp(
                          denom,
                          acts,
                          maxT,
                          maxU,
                          alphabet_size,
                          b,
                          t,
                          U - 1,
                          blank_,
                      )
                      - sigma
                  )
  
                  # now add the weights from big blanks
                  for i in range(num_big_blanks):
                      if t + big_blank_duration[i] < T:
                          # adding to beta[t, U - 1] of weight (in log domain),
                          # beta[t + duration, U - 1] *
                          #     p(big-blank | t, U - 1) / exp(sigma)
                          betas[
                              offset + t * maxU + U - 1
                          ] = rnnt_helper.log_sum_exp(
                              betas[offset + t * maxU + U - 1],
                              betas[
                                  offset
                                  + (t + big_blank_duration[i]) * maxU
                                  + U
                                  - 1
                              ]
                              + logp(
                                  denom,
                                  acts,
                                  maxT,
                                  maxU,
                                  alphabet_size,
                                  b,
                                  t,
                                  U - 1,
                                  blank_ - 1 - i,
                              )
                              - sigma,
                          )
                      elif (
                          t + big_blank_duration[i] == T
                          and big_blank_duration[i] != 1
                      ):
                          # adding to beta[T - duration, U - 1] of weight (in log domain),
                          # p(big-blank | T - duration, U - 1) / exp(sigma)
                          betas[
                              offset + t * maxU + U - 1
                          ] = rnnt_helper.log_sum_exp(
                              betas[offset + t * maxU + U - 1],
                              logp(
                                  denom,
                                  acts,
                                  maxT,
                                  maxU,
                                  alphabet_size,
                                  b,
                                  t,
                                  U - 1,
                                  blank_ - 1 - i,
                              )
                              - sigma,
                          )
  
          elif u < U:
              if t == T - 1:
                  # for u in reversed(range(U - 1)) step to initialize betas[b, T-1, u]
                  betas[offset + (T - 1) * maxU + u] = (
                      betas[offset + (T - 1) * maxU + u + 1]
                      + logp(
                          denom,
                          acts,
                          maxT,
                          maxU,
                          alphabet_size,
                          b,
                          T - 1,
                          u,
                          labels[u],
                      )
                      - sigma
                  )
              elif (t >= 0) and (t < T - 1):
                  # for t in reversed(range(T - 1)) for u in reversed(range(U - 1))
                  # step to compute betas[b, t, u]
                  no_emit = (
                      betas[offset + (t + 1) * maxU + u]
                      + logp(
                          denom,
                          acts,
                          maxT,
                          maxU,
                          alphabet_size,
                          b,
                          t,
                          u,
                          blank_,
                      )
                      - sigma
                  )
                  emit = (
                      betas[offset + t * maxU + u + 1]
                      + logp(
                          denom,
                          acts,
                          maxT,
                          maxU,
                          alphabet_size,
                          b,
                          t,
                          u,
                          labels[u],
                      )
                      - sigma
                  )
                  betas[
                      offset + t * maxU + u
                  ] = rnnt_helper.log_sum_exp(emit, no_emit)
  
                  # now add the weights from big blanks
                  for i in range(num_big_blanks):
                      if t < T - big_blank_duration[i]:
                          # added weight for the big-blank,
                          # beta[t + duration, u] * p(big-blank | t, u) / exp(sigma)
                          big_blank_no_emit = (
                              betas[
                                  offset
                                  + (t + big_blank_duration[i]) * maxU
                                  + u
                              ]
                              + logp(
                                  denom,
                                  acts,
                                  maxT,
                                  maxU,
                                  alphabet_size,
                                  b,
                                  t,
                                  u,
                                  blank_ - 1 - i,
                              )
                              - sigma
                          )
                          betas[
                              offset + t * maxU + u
                          ] = rnnt_helper.log_sum_exp(
                              betas[offset + t * maxU + u],
                              big_blank_no_emit,
                          )
  
          # sync across all B=b and U=u
          cuda.syncthreads()
  
      # After final sync, betas[b, 0, 0] gives
      # log-likelihood of backward pass.
      if u == 0:
          llBackward[b] = betas[offset]
  
file_path: espnet/espnet2/asr/transducer/rnnt_multi_blank/utils/cuda_utils/gpu_rnnt_kernel.py
imports: import math
  import torch
  from numba import cuda
  from espnet2.asr.transducer.rnnt_multi_blank.utils import rnnt_helper
label_desc: Use image node_id 6 for calling a global function with example usage: compute_multiblank_betas_kernel(acts, denom, sigma, betas, llBackward, xlen, ylen, mlabels, minibatch, maxT, maxU, alphabet_size, blank_, big_blank_duration, num_big_blanks) without return types