class_docstr: Pure Python implementation of TDT loss (https://arxiv.org/pdf/2304.06795.pdf)
class_docstr_tok: ["Pure","Python","implementation","of","TDT","loss","(","https",":","\/\/arxiv.org\/pdf\/2304.06795.pdf",")"]
docstr: This function implements Equation 7 in the TDT paper https://arxiv.org/pdf/2304.06795.pdf,
  Simply put, for each alpha(t, u), it sums over the contribution from all incoming blank arcs and non-blank arcs.
docstr_tok: ["This","function","implements","Equation","7","in","the","TDT","paper","https",":","\/\/arxiv.org\/pdf\/2304.06795.pdf",",","Simply","put",",","for","each","alpha","(","t",",","u",")",",","it","sums","over","the","contribution","from","all","incoming","blank","arcs","and","non-blank","arcs","."]
code: def compute_forward_prob(
      self, acts, duration_acts, labels, act_lens, label_lens
  ):
      """This function implements Equation 7 in the TDT paper https://arxiv.org/pdf/2304.06795.pdf,
      Simply put, for each alpha(t, u), it sums over the contribution from all incoming blank arcs and non-blank arcs.
      """
      B, T, U, _ = acts.shape
  
      log_alpha = torch.zeros(B, T, U)
      log_alpha = log_alpha.cuda()
      for b in range(B):
          for t in range(T):
              for u in range(U):
                  if u == 0:
                      if t == 0:
                          # both t and u are 0, this is the base case for alphas.
                          log_alpha[b, t, u] = 0.0
                      else:
                          # u = 0 and t != 0: only considers blank emissions.
                          log_alpha[b, t, u] = -1000.0
                          for n, l in enumerate(self.durations):
                              if (
                                  t - l >= 0 and l > 0
                              ):  # checking conditions for blank emission, l has to be at least 1
                                  tmp = (
                                      log_alpha[b, t - l, u]
                                      + acts[b, t - l, u, self.blank]
                                      + duration_acts[b, t - l, u, n]
                                  )
                                  log_alpha[b, t, u] = self.logsumexp(
                                      tmp, 1.0 * log_alpha[b, t, u]
                                  )
  
                  else:
                      # u != 0 here, need to consider both blanks and non-blanks.
                      log_alpha[b, t, u] = -1000.0
                      for n, l in enumerate(self.durations):
                          if t - l >= 0:
                              if (
                                  l > 0
                              ):  # for blank emissions. Need to ensure index is not out-of-bound.
                                  tmp = (
                                      log_alpha[b, t - l, u]
                                      + acts[b, t - l, u, self.blank]
                                      + duration_acts[b, t - l, u, n]
                                  )
                                  log_alpha[b, t, u] = self.logsumexp(
                                      tmp, 1.0 * log_alpha[b, t, u]
                                  )
  
                              # non-blank emissions.
                              tmp = (
                                  log_alpha[b, t - l, u - 1]
                                  + acts[
                                      b, t - l, u - 1, labels[b, u - 1]
                                  ]
                                  + duration_acts[b, t - l, u - 1, n]
                              )
                              log_alpha[b, t, u] = self.logsumexp(
                                  tmp, 1.0 * log_alpha[b, t, u]
                              )
  
      log_probs = []
      for b in range(B):
          tt = torch.Tensor([-1000.0]).cuda()[0]
  
          # need to loop over all possible ways that blank with different durations contributes to the final loss.
          for n, l in enumerate(self.durations):
              if act_lens[b] - l >= 0 and l > 0:
                  bb = (
                      log_alpha[b, act_lens[b] - l, label_lens[b]]
                      + acts[
                          b, act_lens[b] - l, label_lens[b], self.blank
                      ]
                      + duration_acts[
                          b, act_lens[b] - l, label_lens[b], n
                      ]
                  )
  
                  tt = self.logsumexp(bb, 1.0 * tt)
  
          log_probs.append(tt)
  
      log_prob = torch.stack(log_probs)
  
      return log_prob, log_alpha
  
imports: from typing import List
  import torch
  from nemo.core.classes import Loss
  from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType
label_desc: Use image node_id 6 for calling the TDTLossPytorch obj's underlying member method code with example usage: obj.compute_forward_prob(acts, duration_acts, labels, act_lens, label_lens) and returns: log_prob,  log_alpha