code: def test_unit_schema(self, device, blank_first):
      vocab_size = 3
      blank_id = 0 if blank_first else vocab_size - 1
      if blank_first:
          labels = [1, 1, 2, 1]
      else:
          labels = [1, 1, 0, 1]
      loss = GraphRnntLoss(blank=blank_id)
      unit_schema = loss.get_unit_schema(
          units_tensor=torch.tensor(
              labels, device=torch.device(device)
          ),
          vocab_size=vocab_size,
      )
  
      etalon_schema_fst: List[List[int]] = []
      for label_i, label in enumerate(labels):
          etalon_schema_fst.append(
              [label_i, label_i + 1, label, label, label_i, 0]
          )  # forward: label
          etalon_schema_fst.append(
              [label_i, label_i, blank_id, blank_id, label_i, 0]
          )  # self-loop: blank
      etalon_schema_fst.append(
          [len(labels), len(labels), blank_id, blank_id, len(labels), 0]
      )
      etalon_schema_fst.append(
          [len(labels), len(labels) + 1, -1, -1, -1, 0]
      )  # transition to final state
      etalon_schema_fst.append([len(labels) + 1])  # final state
      etalon_schema_fst = sorted(
          etalon_schema_fst
      )  # required for k2.Fsa.from_str
      etalon_schema_fst_str = "\n".join(
          [" ".join(map(str, line)) for line in etalon_schema_fst]
      )
      etalon_unit_schema = k2.Fsa.from_str(
          etalon_schema_fst_str,
          aux_label_names=["aux_labels", "unit_positions"],
      )
  
      assert unit_schema.num_arcs == etalon_unit_schema.num_arcs
      assert (
          unit_schema.shape == etalon_unit_schema.shape
      )  # (num_states, None)
      assert k2.is_rand_equivalent(
          unit_schema,
          etalon_unit_schema,
          log_semiring=True,
          treat_epsilons_specially=False,
      ), "Unit schema input labels mismatch"
      assert k2.is_rand_equivalent(
          unit_schema.invert(),
          etalon_unit_schema.invert(),
          log_semiring=True,
          treat_epsilons_specially=False,
      ), "Unit schema output labels mismatch"
  
      # swap aux_labels and unit positions to test unit_positions
      unit_schema.aux_labels, unit_schema.unit_positions = (
          unit_schema.unit_positions,
          unit_schema.aux_labels,
      )
      (
          etalon_unit_schema.aux_labels,
          etalon_unit_schema.unit_positions,
      ) = (
          etalon_unit_schema.unit_positions,
          etalon_unit_schema.aux_labels,
      )
      assert k2.is_rand_equivalent(
          unit_schema.invert(),
          etalon_unit_schema.invert(),
          log_semiring=True,
          treat_epsilons_specially=False,
      ), "Unit schema unit positions mismatch"
  
file_path: NeMo/tests/collections/asr/k2/test_graph_transducer.py
imports: from typing import List
  import numpy
  import pytest
  import torch
  from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_numpy import RNNTLoss
label_desc: Use image node_id 2 for calling the TestGraphRnnt obj's underlying member method code with example usage: obj.test_unit_schema(device, blank_first) without return types