code: def transducer_loss_reference(
      x, label, f_len, y_len, blank_idx, loss_grad
  ):
      def log_sum_exp(a, b):
          if a >= b:
              return a + torch.log(1 + torch.exp(b - a))
          else:
              return b + torch.log(1 + torch.exp(a - b))
  
      def forward_alpha(x, label, f_len, y_len, blank_idx):
          B, T, U, V = x.size()
          acc_t = (
              torch.float32
              if x.dtype in [torch.float16, torch.float32]
              else x.dtype
          )
          alpha = torch.zeros((B, T, U), dtype=acc_t, device=x.device)
          for b in range(B):
              alpha[b, 0, 0] = 0
              for t in range(1, f_len[b]):
                  alpha[b, t, 0] = (
                      alpha[b, t - 1, 0] + x[b, t - 1, 0, blank_idx]
                  )
              for u in range(1, y_len[b] + 1):
                  alpha[b, 0, u] = (
                      alpha[b, 0, u - 1]
                      + x[b, 0, u - 1, label[b, u - 1]]
                  )
              for t in range(1, f_len[b]):
                  for u in range(1, y_len[b] + 1):
                      curr_ = (
                          alpha[b, t - 1, u] + x[b, t - 1, u, blank_idx]
                      )
                      next_ = (
                          alpha[b, t, u - 1]
                          + x[b, t, u - 1, label[b, u - 1]]
                      )
                      alpha[b, t, u] = log_sum_exp(curr_, next_)
          return alpha
  
      def forward_beta(x, label, f_len, y_len, blank_idx):
          B, T, U, V = x.shape
          acc_t = (
              torch.float32
              if x.dtype in [torch.float16, torch.float32]
              else x.dtype
          )
          beta = torch.zeros((B, T, U), dtype=acc_t, device=x.device)
          for b in range(B):
              beta[b, f_len[b] - 1, y_len[b]] = x[
                  b, f_len[b] - 1, y_len[b], blank_idx
              ]
              for t in range(f_len[b] - 2, -1, -1):
                  beta[b, t, y_len[b]] = (
                      beta[b, t + 1, y_len[b]]
                      + x[b, t, y_len[b], blank_idx]
                  )
              for u in range(y_len[b] - 1, -1, -1):
                  beta[b, f_len[b] - 1, u] = (
                      beta[b, f_len[b] - 1, u + 1]
                      + x[b, f_len[b] - 1, u, label[b, u]]
                  )
              for t in range(f_len[b] - 2, -1, -1):
                  for u in range(y_len[b] - 1, -1, -1):
                      curr_ = beta[b, t + 1, u] + x[b, t, u, blank_idx]
                      next_ = (
                          beta[b, t, u + 1] + x[b, t, u, label[b, u]]
                      )
                      beta[b, t, u] = log_sum_exp(curr_, next_)
          return beta
  
      def backward(
          x, label, f_len, y_len, alpha, beta, loss_grad, blank_idx
      ):
          grad = torch.zeros_like(x)
          B, T, U, V = x.size()
          for b in range(B):
              common_factor = (
                  torch.log(loss_grad[b]) + alpha - beta[b, 0, 0]
              )
              # next
              for u in range(y_len[b]):
                  grad[b, : f_len[b], u, label[b, u]] = -torch.exp(
                      common_factor[b, : f_len[b], u]
                      + beta[b, : f_len[b], u + 1]
                      + x[b, : f_len[b], u, label[b, u]]
                  )
  
              # current
              grad[
                  b, : f_len[b] - 1, : y_len[b] + 1, blank_idx
              ] = -torch.exp(
                  common_factor[b, : f_len[b] - 1, : y_len[b] + 1]
                  + beta[b, 1 : f_len[b], : y_len[b] + 1]
                  + x[b, : f_len[b] - 1, : y_len[b] + 1, blank_idx]
              )
  
              grad[b, f_len[b] - 1, y_len[b], blank_idx] = -torch.exp(
                  common_factor[b, f_len[b] - 1, y_len[b]]
                  + x[b, f_len[b] - 1, y_len[b], blank_idx]
              )
  
          return grad
  
      x_log = torch.nn.functional.log_softmax(x, dim=-1)
      alpha = forward_alpha(x_log, label, f_len, y_len, blank_idx)
      beta = forward_beta(x_log, label, f_len, y_len, blank_idx)
      grad = backward(
          x_log, label, f_len, y_len, alpha, beta, loss_grad, blank_idx
      )
      x_log.backward(grad)
      loss = -beta[:, 0, 0]
      loss = loss.to(x.dtype)
      return alpha, beta, x.grad, loss
  
label_desc: Use image node_id 1 for calling a global function with example usage: transducer_loss_reference(x, label, f_len, y_len, blank_idx, loss_grad) and returns: alpha,  beta,  x,  loss, alpha, beta, grad, unknown, unknown