class_docstr: RNN-T loss implementation based on WFST according
  to "Powerful and Extensible WFST Framework for RNN-Transducer Losses"
  https://ieeexplore.ieee.org/document/10096679
docstr: Construct the RNN-T lattice directly (Grid-Transducer).
  
  Args:
      units_tensor: 1d tensor with text units
      num_frames: length of the sequence (number of frames)
      vocab_size: number of total labels (vocab size including blank)
  
  Returns:
      transducer lattice (k2.Fsa).
      Labels: <unit>:<frame_index>:<unit_position> (k2.Fsa: labels, aux_labels, unit_positions)
docstr_tok: ["Construct","the","RNN-T","lattice","directly","(","Grid-Transducer",")",".","Args",":","units_tensor",":","1d","tensor","with","text","units","num_frames",":","length","of","the","sequence","(","number","of","frames",")","vocab_size",":","number","of","total","labels","(","vocab","size","including","blank",")","Returns",":","transducer","lattice","(","k2.Fsa",")",".","Labels",":","<","unit",">",":","<","frame_index",">",":","<","unit_position",">","(","k2.Fsa",":","labels",",","aux_labels",",","unit_positions",")"]
code: def get_grid(
      self, units_tensor: torch.Tensor, num_frames: int, vocab_size: int
  ) -> "k2.Fsa":
      """
      Construct the RNN-T lattice directly (Grid-Transducer).
  
      Args:
          units_tensor: 1d tensor with text units
          num_frames: length of the sequence (number of frames)
          vocab_size: number of total labels (vocab size including blank)
  
      Returns:
          transducer lattice (k2.Fsa).
          Labels: <unit>:<frame_index>:<unit_position> (k2.Fsa: labels, aux_labels, unit_positions)
      """
      blank_id = self.blank
      text_length = units_tensor.shape[0]
      device = units_tensor.device
      num_grid_states = num_frames * (text_length + 1)
      num_forward_arcs = (num_frames - 1) * (text_length + 1)
      num_text_arcs = text_length * num_frames
      arcs = torch.zeros(
          (num_forward_arcs + num_text_arcs + 2, 4),
          dtype=torch.int32,
          device=device,
      )
      # blank transitions
      # i, i+<text_len + 1>, 0 <blank>, i / <text_len+1>, i % <text_len + 1>
      from_states = torch.arange(num_forward_arcs, device=device)
      to_states = from_states + (text_length + 1)
      arcs[:num_forward_arcs, 0] = from_states
      arcs[:num_forward_arcs, 1] = to_states
      arcs[:num_forward_arcs, 2] = blank_id
  
      # text arcs
      from_states = (
          torch.arange(
              num_grid_states, dtype=torch.int32, device=device
          )
          .reshape(num_frames, text_length + 1)[:, :-1]
          .flatten()
      )
      to_states = from_states + 1
      ilabels = units_tensor.expand(num_frames, -1).flatten()
      arcs[num_forward_arcs:-2, 0] = from_states
      arcs[num_forward_arcs:-2, 1] = to_states
      arcs[num_forward_arcs:-2, 2] = ilabels
  
      # last 2 states
      arcs[-2, :3] = torch.tensor(
          (num_grid_states - 1, num_grid_states, blank_id),
          dtype=torch.int32,
          device=device,
      )
      arcs[-1, :3] = torch.tensor(
          (num_grid_states, num_grid_states + 1, -1),
          dtype=torch.int32,
          device=device,
      )
  
      # sequence indices, time indices
      olabels = torch.div(
          arcs[:, 0], (text_length + 1), rounding_mode="floor"
      )  # arcs[:, 0] // (text_length + 1)
      unit_positions = arcs[:, 0] % (text_length + 1)
      # last state: final
      olabels[-1] = -1
      unit_positions[-1] = -1
  
      # relabel
      # instead of using top sort (extremely expensive) k2.top_sort(rnnt_graph)
      arcs[:-2, 0] = self.relabel_states(
          arcs[:-2, 0], text_length + 1, num_frames
      )
      arcs[:-3, 1] = self.relabel_states(
          arcs[:-3, 1], text_length + 1, num_frames
      )
  
      # sort by start state - required in k2
      # TODO: maybe it is more optimal to avoid sort, construct arcs in ascending order
      _, indices = torch.sort(arcs[:, 0], dim=0)
      sorted_arcs = arcs[indices]
      olabels = olabels[indices]
      unit_positions = unit_positions[indices]
  
      rnnt_graph = k2.Fsa(sorted_arcs, olabels)
      rnnt_graph.unit_positions = unit_positions
      return rnnt_graph
  
imports: import abc
  from contextlib import nullcontext
  from typing import ContextManager
  import torch
  import torch.nn.functional
  from nemo.core.classes.loss import Loss
  from nemo.core.utils.k2_guard import k2
label_desc: Use image node_id 5 for calling the GraphRnntLoss obj's underlying member method code with example usage: obj.get_grid(units_tensor, num_frames, vocab_size) and returns: rnnt_graph