code: def test_grid_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]
      text_length = len(labels)
      num_frames = 5
      loss = GraphRnntLoss(blank=blank_id)
      grid_schema = loss.get_grid(
          units_tensor=torch.tensor(
              labels, device=torch.device(device)
          ),
          num_frames=num_frames,
          vocab_size=vocab_size,
      )
  
      etalon_schema_fst: List[List[int]] = []
      for frame_i in range(num_frames):
          for label_i in range(text_length + 1):
              state = frame_i * (text_length + 1) + label_i
              if label_i < text_length:
                  next_state_label = state + 1
                  # next unit
                  etalon_schema_fst.append(
                      [
                          state,
                          next_state_label,
                          labels[label_i],
                          frame_i,
                          label_i,
                          0,
                      ]
                  )
              if frame_i < num_frames - 1:
                  next_state_frame = (frame_i + 1) * (
                      text_length + 1
                  ) + label_i
                  # next time frame (blank)
                  etalon_schema_fst.append(
                      [
                          state,
                          next_state_frame,
                          blank_id,
                          frame_i,
                          label_i,
                          0,
                      ]
                  )
  
      last_grid_state = num_frames * (text_length + 1) - 1
      etalon_schema_fst.append(
          [
              last_grid_state,
              last_grid_state + 1,
              blank_id,
              num_frames - 1,
              text_length,
              0,
          ]
      )
      etalon_schema_fst.append(
          [last_grid_state + 1, last_grid_state + 2, -1, -1, -1, 0]
      )  # transition to final state
      etalon_schema_fst.append([last_grid_state + 2])  # 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_grid_schema = k2.Fsa.from_str(
          etalon_schema_fst_str,
          aux_label_names=["aux_labels", "unit_positions"],
      )
  
      assert grid_schema.num_arcs == etalon_grid_schema.num_arcs
      assert (
          grid_schema.shape == etalon_grid_schema.shape
      )  # (num_states, None)
      assert k2.is_rand_equivalent(
          grid_schema,
          etalon_grid_schema,
          log_semiring=True,
          treat_epsilons_specially=False,
      ), "Grid schema input labels mismatch"
      assert k2.is_rand_equivalent(
          grid_schema.invert(),
          etalon_grid_schema.invert(),
          log_semiring=True,
          treat_epsilons_specially=False,
      ), "Grid schema output labels mismatch"
  
      # swap aux_labels and unit positions to test unit_positions
      grid_schema.aux_labels, grid_schema.unit_positions = (
          grid_schema.unit_positions,
          grid_schema.aux_labels,
      )
      (
          etalon_grid_schema.aux_labels,
          etalon_grid_schema.unit_positions,
      ) = (
          etalon_grid_schema.unit_positions,
          etalon_grid_schema.aux_labels,
      )
      assert k2.is_rand_equivalent(
          grid_schema.invert(),
          etalon_grid_schema.invert(),
          log_semiring=True,
          treat_epsilons_specially=False,
      ), "Grid 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 3 for calling the TestGraphRnnt obj's underlying member method code with example usage: obj.test_grid_schema(device, blank_first) without return types