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: Get unit schema (target text) graph for RNN-T loss (Compose-Transducer).
  Forward arcs represent text labels.
  
  Example graph: text [1, 2], blank=0.
  
  graph::
  
          0:0:0                  0:0:1                  0:0:2
        +-------+              +-------+              +-------+
        v       |              v       |              v       |
      +-----------+  1:1:0   +-----------+  2:2:1   +-----------+  -1:-1:-1  #===#
      |     0     | -------> |     1     | -------> |     2     | ---------> H 3 H
      +-----------+          +-----------+          +-----------+            #===#
  
  Args:
      units_tensor: 1d tensor with text units
      vocab_size: number of total labels (vocab size including blank)
  
  Returns:
      unit schema graph (k2.Fsa).
      Labels: <unit>:<unit>:<unit_position> (k2.Fsa: labels, aux_labels, unit_positions)
docstr_tok: ["Get","unit","schema","(","target","text",")","graph","for","RNN-T","loss","(","Compose-Transducer",")",".","Forward","arcs","represent","text","labels",".","Example","graph",":","text","[","1",",","2","]",",","blank=0",".","graph",":",":","0:0:0","0:0:1","0:0:2","+","--","--","--","-+","+","--","--","--","-+","+","--","--","--","-+","v","|","v","|","v","|","+","--","--","--","--","--","-+","1:1:0","+","--","--","--","--","--","-+","2:2:1","+","--","--","--","--","--","-+","-1",":","-1",":","-1","#","===","#","|","0","|","--","--","--","-",">","|","1","|","--","--","--","-",">","|","2","|","--","--","--","--","-",">","H","3","H","+","--","--","--","--","--","-+","+","--","--","--","--","--","-+","+","--","--","--","--","--","-+","#","===","#","Args",":","units_tensor",":","1d","tensor","with","text","units","vocab_size",":","number","of","total","labels","(","vocab","size","including","blank",")","Returns",":","unit","schema","graph","(","k2.Fsa",")",".","Labels",":","<","unit",">",":","<","unit",">",":","<","unit_position",">","(","k2.Fsa",":","labels",",","aux_labels",",","unit_positions",")"]
code: def get_unit_schema(
      self, units_tensor: torch.Tensor, vocab_size: int
  ) -> "k2.Fsa":
      """
      Get unit schema (target text) graph for RNN-T loss (Compose-Transducer).
      Forward arcs represent text labels.
  
      Example graph: text [1, 2], blank=0.
  
      graph::
  
              0:0:0                  0:0:1                  0:0:2
            +-------+              +-------+              +-------+
            v       |              v       |              v       |
          +-----------+  1:1:0   +-----------+  2:2:1   +-----------+  -1:-1:-1  #===#
          |     0     | -------> |     1     | -------> |     2     | ---------> H 3 H
          +-----------+          +-----------+          +-----------+            #===#
  
      Args:
          units_tensor: 1d tensor with text units
          vocab_size: number of total labels (vocab size including blank)
  
      Returns:
          unit schema graph (k2.Fsa).
          Labels: <unit>:<unit>:<unit_position> (k2.Fsa: labels, aux_labels, unit_positions)
      """
  
      blank_id = self.blank
      device = units_tensor.device
      text_len = units_tensor.shape[0]
  
      # arcs
      # text_len + 1 states, in every state - self-loops (blank) and forward (text label / last forward -1)
      arcs = torch.zeros(
          ((text_len + 1) * 2, 4), dtype=torch.int32, device=device
      )
      text_indices = torch.arange(
          0, text_len + 1, dtype=torch.int32, device=device
      )
      # blank labels
      arcs[::2, 0] = text_indices  # from state
      arcs[::2, 1] = text_indices  # to state
      arcs[::2, 2] = blank_id
  
      # text labels
      arcs[1::2, 0] = text_indices  # from state
      arcs[1::2, 1] = text_indices + 1  # to state
      arcs[1:-1:2, 2] = units_tensor  # labels: text
  
      arcs[
          -1, 2
      ] = (
          -1
      )  # last transition to final state, ilabel=-1 (special for k2)
      olabels = arcs[:, 2].detach().clone()  # same as ilabels
  
      fsa_text = k2.Fsa(arcs, olabels)
      fsa_text.unit_positions = (
          text_indices.expand(2, -1).transpose(0, 1).flatten()
      )
      fsa_text.unit_positions[-1] = -1  # last transition to final state
      return fsa_text
  
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 2 for calling the GraphRnntLoss obj's underlying member method code with example usage: obj.get_unit_schema(units_tensor, vocab_size) and returns: fsa_text