class_docstr: Pure Python implementation of multi-blank transducer loss (https://arxiv.org/pdf/2211.03541.pdf)
class_docstr_tok: ["Pure","Python","implementation","of","multi-blank","transducer","loss","(","https",":","\/\/arxiv.org\/pdf\/2211.03541.pdf",")"]
code: def compute_forward_prob(self, acts, labels, act_lens, label_lens):
      B, T, U, _ = acts.shape
  
      log_alpha = torch.zeros(B, T, U, device=acts.device)
      for t in range(T):
          for u in range(U):
              if u == 0:
                  if t == 0:
                      # this is the base case: (t=0, u=0) with log-alpha = 0.
                      log_alpha[:, t, u] = 0.0
                  else:
                      # this is case for (t = 0, u > 0), reached by (t, u - d)
                      # emitting a blank symbol of duration d.
                      log_alpha[:, t, u] = (
                          log_alpha[:, t - 1, u]
                          + acts[:, t - 1, 0, self.blank]
                      )
                      for i, d in enumerate(self.big_blank_durations):
                          if t >= d:
                              tt = (
                                  log_alpha[:, t - d, u]
                                  + acts[
                                      :, t - d, 0, self.blank - 1 - i
                                  ]
                              )
                              log_alpha[:, t, u] = torch.logsumexp(
                                  torch.stack(
                                      [1.0 * log_alpha[:, t, u], tt]
                                  ),
                                  dim=0,
                              )
  
              else:
                  if t == 0:
                      # in case of (u > 0, t = 0), this is only reached from
                      # (t, u - 1) with a label emission.
                      gathered = torch.gather(
                          acts[:, t, u - 1],
                          dim=1,
                          index=labels[:, u - 1]
                          .view(-1, 1)
                          .type(torch.int64),
                      ).reshape(-1)
                      log_alpha[:, t, u] = (
                          log_alpha[:, t, u - 1] + gathered
                      )
                  else:
                      # here both t and u are > 0, this state is reachable
                      # with two possibilities: (t - d, u) with emission of
                      # blank with duration d, or (t, u - 1) with a label emission.
  
                      # first we take care of the standard blank.
                      log_alpha[:, t, u] = torch.logsumexp(
                          torch.stack(
                              [
                                  log_alpha[:, t - 1, u]
                                  + acts[:, t - 1, u, self.blank],
                                  log_alpha[:, t, u - 1]
                                  + torch.gather(
                                      acts[:, t, u - 1],
                                      dim=1,
                                      index=labels[:, u - 1]
                                      .view(-1, 1)
                                      .type(torch.int64),
                                  ).reshape(-1),
                              ]
                          ),
                          dim=0,
                      )
  
                      # now we go over all big blanks. They need to be considered if current t >= blank duration d.
                      for i, d in enumerate(self.big_blank_durations):
                          if t >= d:
                              tt = (
                                  log_alpha[:, t - d, u]
                                  + acts[
                                      :, t - d, u, self.blank - 1 - i
                                  ]
                              )
                              log_alpha[:, t, u] = torch.logsumexp(
                                  torch.stack(
                                      [1.0 * log_alpha[:, t, u], tt]
                                  ),
                                  dim=0,
                              )
  
      log_probs = []
      for b in range(B):
          # here we need to add the final blank emission weights, which needs
          # to consider all possible blank durations.
          to_append = (
              log_alpha[b, act_lens[b] - 1, label_lens[b]]
              + acts[b, act_lens[b] - 1, label_lens[b], self.blank]
          )
  
          for i, d in enumerate(self.big_blank_durations):
              if act_lens[b] >= d:
                  tt = (
                      log_alpha[b, act_lens[b] - d, label_lens[b]]
                      + acts[
                          b,
                          act_lens[b] - d,
                          label_lens[b],
                          self.blank - 1 - i,
                      ]
                  )
                  to_append = torch.logsumexp(
                      torch.stack([1.0 * to_append, tt]), dim=0
                  )
  
          log_probs.append(to_append)
      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 5 for calling the MultiblankRNNTLossPytorch obj's underlying member method code with example usage: obj.compute_forward_prob(acts, labels, act_lens, label_lens) and returns: log_prob,  log_alpha