class_docstr: W-Transducer loss: RNN-T loss modification for training RNN-T model for the case
  when some text at the beginning/end of the utterance is missing.
  The resulting model behaves like the RNN-T model (no modification for decoding is required).
  For details see "Powerful and Extensible WFST Framework for RNN-Transducer Losses" paper
      https://ieeexplore.ieee.org/document/10096679
class_docstr_tok: ["W-Transducer","loss",":","RNN-T","loss","modification","for","training","RNN-T","model","for","the","case","when","some","text","at","the","beginning\/end","of","the","utterance","is","missing",".","The","resulting","model","behaves","like","the","RNN-T","model","(","no","modification","for","decoding","is","required",")",".","For","details","see","``","Powerful","and","Extensible","WFST","Framework","for","RNN-Transducer","Losses","''","paper","https",":","\/\/ieeexplore.ieee.org\/document\/10096679"]
docstr: Get temporal schema graph for W-Transducer loss (Compose-Transducer).
  
  Example graph: blank=0, num_frames=3, vocab_size=3, last_blank_mode="force_final".
  Labels: <unit>:<frame_index>. <unit> is a unit from vocab + special eps ids `vocab_size`, `vocab_size+1`.
  
  graph for force_final::
  
                                                   4:0
                 +--------------------------------------------+
                 |                               4:1          |
                 |                     +--------------------+ |
          1:0    |              1:1    |              1:2   | |
        +-----+  |            +-----+  |            +-----+ | |
        v     |  |            v     |  |            v     | v v
      +--------------+  0:0  +------------+  0:1   +------------+  0:2   +---+  -1:-1   #===#
      |    0         | ----> |    1       | -----> |    2       | -----> | 3 | -------> H 4 H
      +--------------+       +------------+        +------------+        +---+          #===#
        ^ 2:0 |  |  |         ^ 2:1 |  ^            ^ 2:2 |  ^
        +-----+  |  |         +-----+  |            +-----+  |
                 |  |     3:0          |                     |
                 |  +------------------+     3:0             |
                 +-------------------------------------------+
  
  
  Args:
      num_frames: length of the sequence (in frames)
      vocab_size: number of labels (including blank)
      device: device for tensor to construct
  
  Returns:
      temporal schema graph (k2.Fsa).
      Labels: <unit>:<frame_index>. <unit> is a unit from vocab + special units (e.g., additional eps).
docstr_tok: ["Get","temporal","schema","graph","for","W-Transducer","loss","(","Compose-Transducer",")",".","Example","graph",":","blank=0",",","num_frames=3",",","vocab_size=3",",","last_blank_mode=","''","force_final","''",".","Labels",":","<","unit",">",":","<","frame_index",">",".","<","unit",">","is","a","unit","from","vocab","+","special","eps","ids","`","vocab_size","`",",","`","vocab_size+1","`",".","graph","for","force_final",":",":","4:0","+","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","+","|","4:1","|","|","+","--","--","--","--","--","--","--","--","--","--","+","|","1:0","|","1:1","|","1:2","|","|","+","--","--","-+","|","+","--","--","-+","|","+","--","--","-+","|","|","v","|","|","v","|","|","v","|","v","v","+","--","--","--","--","--","--","--","+","0:0","+","--","--","--","--","--","--","+","0:1","+","--","--","--","--","--","--","+","0:2","+","--","-+","-1",":","-1","#","===","#","|","0","|","--","--",">","|","1","|","--","--","-",">","|","2","|","--","--","-",">","|","3","|","--","--","--","-",">","H","4","H","+","--","--","--","--","--","--","--","+","+","--","--","--","--","--","--","+","+","--","--","--","--","--","--","+","+","--","-+","#","===","#","^","2:0","|","|","|","^","2:1","|","^","^","2:2","|","^","+","--","--","-+","|","|","+","--","--","-+","|","+","--","--","-+","|","|","|","3:0","|","|","|","+","--","--","--","--","--","--","--","--","--","+","3:0","|","+","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","-+","Args",":","num_frames",":","length","of","the","sequence","(","in","frames",")","vocab_size",":","number","of","labels","(","including","blank",")","device",":","device","for","tensor","to","construct","Returns",":","temporal","schema","graph","(","k2.Fsa",")",".","Labels",":","<","unit",">",":","<","frame_index",">",".","<","unit",">","is","a","unit","from","vocab","+","special","units","(","e.g.",",","additional","eps",")","."]
code: def get_temporal_schema(
      self, num_frames: int, vocab_size: int, device: torch.device
  ) -> "k2.Fsa":
      """
      Get temporal schema graph for W-Transducer loss (Compose-Transducer).
  
      Example graph: blank=0, num_frames=3, vocab_size=3, last_blank_mode="force_final".
      Labels: <unit>:<frame_index>. <unit> is a unit from vocab + special eps ids `vocab_size`, `vocab_size+1`.
  
      graph for force_final::
  
                                                       4:0
                     +--------------------------------------------+
                     |                               4:1          |
                     |                     +--------------------+ |
              1:0    |              1:1    |              1:2   | |
            +-----+  |            +-----+  |            +-----+ | |
            v     |  |            v     |  |            v     | v v
          +--------------+  0:0  +------------+  0:1   +------------+  0:2   +---+  -1:-1   #===#
          |    0         | ----> |    1       | -----> |    2       | -----> | 3 | -------> H 4 H
          +--------------+       +------------+        +------------+        +---+          #===#
            ^ 2:0 |  |  |         ^ 2:1 |  ^            ^ 2:2 |  ^
            +-----+  |  |         +-----+  |            +-----+  |
                     |  |     3:0          |                     |
                     |  +------------------+     3:0             |
                     +-------------------------------------------+
  
  
      Args:
          num_frames: length of the sequence (in frames)
          vocab_size: number of labels (including blank)
          device: device for tensor to construct
  
      Returns:
          temporal schema graph (k2.Fsa).
          Labels: <unit>:<frame_index>. <unit> is a unit from vocab + special units (e.g., additional eps).
      """
      blank_id = self.blank
      start_eps_id = vocab_size
      end_eps_id = vocab_size + 1
      num_eps = 2
  
      num_sequence_arcs = (
          num_frames * vocab_size + (num_frames - 1) * num_eps + 1
      )
      fsa_temporal_arcs = torch.zeros(
          (num_sequence_arcs, 4), dtype=torch.int32, device=device
      )
      sequence_states = torch.arange(
          0, num_frames, dtype=torch.int32, device=device
      )
      sequence_states_next = sequence_states + 1
      # for every state - vocab_size+1 arcs, [0, 1, ..., vocab_size-1, eps, 0, 1, ..., vocab_size-1, eps, ...]
      start_states = (
          sequence_states.expand(vocab_size + num_eps, num_frames)
          .transpose(0, 1)
          .flatten()
      )
  
      # self-loops - all, make forward arcs later
      fsa_temporal_arcs[:num_sequence_arcs, 0] = start_states[
          :-1
      ]  # from
      fsa_temporal_arcs[:num_sequence_arcs, 1] = start_states[:-1]  # to
      fsa_temporal_arcs[:num_sequence_arcs, 2] = (
          torch.arange(
              0, vocab_size + num_eps, dtype=torch.int32, device=device
          )
          .expand(num_frames, vocab_size + num_eps)
          .flatten()[:-1]
      )
      # forward arcs
      fsa_temporal_arcs[
          blank_id : num_sequence_arcs : vocab_size + num_eps, 1
      ] = sequence_states_next  # blanks
      # eps arcs
      fsa_temporal_arcs[
          start_eps_id : num_sequence_arcs : vocab_size + num_eps, 0
      ] = 0
      fsa_temporal_arcs[
          start_eps_id : num_sequence_arcs : vocab_size + num_eps, 1
      ] = (sequence_states + 1)
      fsa_temporal_arcs[
          end_eps_id : num_sequence_arcs : vocab_size + num_eps, 0
      ] = sequence_states[:-1]
      fsa_temporal_arcs[
          end_eps_id : num_sequence_arcs : vocab_size + num_eps, 1
      ] = (
          num_frames - 1
          if self.last_blank_mode == self.LastBlankMode.FORCE_FINAL
          else num_frames
      )
  
      # transition to last final state
      fsa_temporal_arcs[-1, :3] = torch.tensor(
          (num_frames, num_frames + 1, -1),
          dtype=torch.int32,
          device=device,
      )
  
      # need to sort arcs
      _, indices = torch.sort(fsa_temporal_arcs[:, 0], dim=0)
      fsa_temporal_arcs = fsa_temporal_arcs[indices]
  
      # output symbols: position in the sequence, same as start states for arcs
      olabels = fsa_temporal_arcs[:, 0].detach().clone()
      olabels[-1] = -1  # transition to the last final state
  
      fsa_temporal = k2.Fsa(fsa_temporal_arcs, olabels)
      fsa_temporal = k2.arc_sort(fsa_temporal)  # need for compose
      return fsa_temporal
  
file_path: NeMo/nemo/collections/asr/parts/k2/w_transducer.py
imports: from contextlib import nullcontext
  from typing import Union
  import torch
  import torch.nn.functional
  from nemo.collections.asr.parts.k2.graph_transducer import GraphRnntLoss, force_float32_context
  from nemo.core.utils.k2_guard import k2
  from nemo.utils.enum import PrettyStrEnum
label_desc: Use image node_id 3 for calling the GraphWTransducerLoss obj's underlying member method code with example usage: obj.get_temporal_schema(num_frames, vocab_size, device) and returns: fsa_temporal