class_docstr: k2 Mixin class that simplifies the construction of various k2-based RNNT-like losses. Inherits CtcK2Mixin.
  
  It does the following:
      -   Prepares and adapts the input tensors.
      -   Creates Emissions graphs.
      -   Extracts the labels and probabilities of the best lattice path (method _extract_labels_and_probabilities).
class_docstr_tok: ["k2","Mixin","class","that","simplifies","the","construction","of","various","k2-based","RNNT-like","losses",".","Inherits","CtcK2Mixin",".","It","does","the","following",":","-","Prepares","and","adapts","the","input","tensors",".","-","Creates","Emissions","graphs",".","-","Extracts","the","labels","and","probabilities","of","the","best","lattice","path","(","method","_extract_labels_and_probabilities",")","."]
docstr: Before calling super()._prepare_log_probs_and_targets, this method reshapes the log_probs tensor 
  from (B, T, U+1, D) to (B, T', D) where T' = T*(U+1), shifts paddings along T and U towards the end of T', 
  and recomputes input_lengths.
  
  It also calculates indices on which <epsilon> steps should be applied to the log_probs tensor to emulate 
  <blank> arcs shift of the Emissions graph for the pruned RNNT variant.
docstr_tok: ["Before","calling","super","(",")","._prepare_log_probs_and_targets",",","this","method","reshapes","the","log_probs","tensor","from","(","B",",","T",",","U+1",",","D",")","to","(","B",",","T","'",",","D",")","where","T","'","=","T","*","(","U+1",")",",","shifts","paddings","along","T","and","U","towards","the","end","of","T","'",",","and","recomputes","input_lengths",".","It","also","calculates","indices","on","which","<","epsilon",">","steps","should","be","applied","to","the","log_probs","tensor","to","emulate","<","blank",">","arcs","shift","of","the","Emissions","graph","for","the","pruned","RNNT","variant","."]
code: def _prepare_log_probs_and_targets(
      self,
      log_probs: torch.Tensor,
      input_lengths: torch.Tensor,
      targets: torch.Tensor,
      target_lengths: torch.Tensor,
  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
      """Before calling super()._prepare_log_probs_and_targets, this method reshapes the log_probs tensor
      from (B, T, U+1, D) to (B, T', D) where T' = T*(U+1), shifts paddings along T and U towards the end of T',
      and recomputes input_lengths.
  
      It also calculates indices on which <epsilon> steps should be applied to the log_probs tensor to emulate
      <blank> arcs shift of the Emissions graph for the pruned RNNT variant.
      """
      assert len(log_probs.size()) == 4  # B T U D
      B, T, U, D = log_probs.size()
      TU = T * U
  
      # save step indices if, as we assume, decoder output pruning has been applied
      if (
          self.predictor_window_size > 0
          and self.predictor_window_size < target_lengths.max()
      ):
          window_size_with_blank = self.predictor_window_size + 1
          ranges_begin = get_uniform_rnnt_prune_ranges(
              input_lengths,
              target_lengths,
              window_size_with_blank,
              self.predictor_step_size,
              T,
              True,
          )
          step_sizes = ranges_begin[:, 1:] - ranges_begin[:, :-1]
          raw_step_indices = torch.where(step_sizes > 0)
          if self.predictor_step_size > 1:
              raw_step_indices = torch.repeat_interleave(
                  torch.stack(raw_step_indices).T,
                  step_sizes[raw_step_indices],
                  dim=0,
              ).T
              raw_step_indices = (
                  raw_step_indices[0],
                  raw_step_indices[1],
              )
          unique, count = torch.unique(
              raw_step_indices[0], return_counts=True
          )
          shift_mask = raw_step_indices[0].unsqueeze(0).repeat(
              len(unique), 1
          ) == unique.unsqueeze(-1)
          step_indices = (
              raw_step_indices[0],
              (
                  torch.arange(ranges_begin.size(1))
                  .unsqueeze(0)
                  .repeat(ranges_begin.size(0), 1)
                  * window_size_with_blank
              )[(raw_step_indices[0], raw_step_indices[1] + 1)]
              + torch.cumsum(shift_mask, 1)[shift_mask]
              - 1,
          )
          max_count = count.max()
          max_count_vec = torch.full((B,), max_count)
          max_count_vec[unique] -= count
          pad_indices_row = torch.repeat_interleave(
              torch.arange(B), max_count_vec
          )
          pad_unique = torch.unique(pad_indices_row)
          pad_shift_mask = pad_indices_row.unsqueeze(0).repeat(
              len(pad_unique), 1
          ) == pad_unique.unsqueeze(-1)
          pad_indices = (
              pad_indices_row,
              T * window_size_with_blank
              + max_count
              - torch.cumsum(pad_shift_mask, 1)[pad_shift_mask],
          )
          self.__step_indices = (
              torch.cat((step_indices[0], pad_indices[0])),
              torch.cat((step_indices[1], pad_indices[1])),
          )
          self.__supervisions_add = max_count - max_count_vec
      else:
          self.__step_indices = None
          self.__supervisions_add = None
  
      # reshape 4D log_probs to 3D with respect to target_lengths
      non_pad_mask_true = make_non_pad_mask_3d(
          input_lengths, target_lengths + 1, T, U
      ).flatten(1)
      input_lengths = non_pad_mask_true.sum(1)
      non_pad_mask_fake = make_non_pad_mask(input_lengths, TU).flatten()
      non_pad_mask_true = non_pad_mask_true.flatten()
      rearranged_indices = torch.arange(TU * B, device=log_probs.device)
      rearranged_indices_buffer = rearranged_indices.clone()
      rearranged_indices[non_pad_mask_fake] = rearranged_indices_buffer[
          non_pad_mask_true
      ]
      rearranged_indices[
          ~non_pad_mask_fake
      ] = rearranged_indices_buffer[~non_pad_mask_true]
      log_probs = log_probs.reshape(-1, D)[rearranged_indices].view(
          B, -1, D
      )
  
      return super()._prepare_log_probs_and_targets(
          log_probs, input_lengths, targets, target_lengths
      )
  
file_path: NeMo/nemo/collections/asr/parts/k2/loss_mixins.py
imports: from abc import ABC
  from typing import List, Optional, Tuple
  import torch
  from nemo.collections.asr.parts.k2.grad_utils import GradExpNormalize
  from nemo.collections.asr.parts.k2.utils import create_supervision, get_arc_weights, get_uniform_rnnt_prune_ranges, make_non_pad_mask, make_non_pad_mask_3d, prep_padded_densefsavec
  from nemo.core.utils.k2_guard import k2
label_desc: Use image node_id 1 for calling the RnntK2Mixin obj's underlying member method code with example usage: obj._prepare_log_probs_and_targets(log_probs, input_lengths, targets, target_lengths) and returns: super