class_docstr: Base class for graph transducer losses.
  Implementation of the approach described in "Powerful and Extensible WFST Framework for RNN-Transducer Losses"
  https://ieeexplore.ieee.org/document/10096679
  
  Compose-Transducer: compose the unit (target text) and temporal schemas (graphs) into lattice.
      Subclass should implement `get_unit_schema` and `get_temporal_schema` methods.
  Grid-Transducer: construct the RNN-T lattice (grid) directly in code.
      Subclass should implement `get_grid` method.
docstr: Get indices of flatten logits for each arc in the lattices.
  
  Args:
      target_fsas_vec: batch of target FSAs with lattices
      logits_shape: shape of the logits tensor
  
  Returns:
      1d tensor with indices
docstr_tok: ["Get","indices","of","flatten","logits","for","each","arc","in","the","lattices",".","Args",":","target_fsas_vec",":","batch","of","target","FSAs","with","lattices","logits_shape",":","shape","of","the","logits","tensor","Returns",":","1d","tensor","with","indices"]
code: def get_logits_indices(
      self, target_fsas_vec: k2.Fsa, logits_shape: torch.Size
  ) -> torch.Tensor:
      """
      Get indices of flatten logits for each arc in the lattices.
  
      Args:
          target_fsas_vec: batch of target FSAs with lattices
          logits_shape: shape of the logits tensor
  
      Returns:
          1d tensor with indices
      """
      # logits_shape: B x Time x Text+1 x Labels
      batch_size = logits_shape[0]
      device = target_fsas_vec.device
      scores_to_batch_i = torch.repeat_interleave(
          torch.arange(batch_size, device=device, dtype=torch.int64),
          torch.tensor(
              [
                  target_fsas_vec.arcs.index(0, i)[0].values().shape[0]
                  for i in range(batch_size)
              ],
              device=device,
          ),
      )
      indices = (
          scores_to_batch_i
          * logits_shape[1]
          * logits_shape[2]
          * logits_shape[3]  # Batch
          + target_fsas_vec.aux_labels.to(torch.int64)
          * logits_shape[2]
          * logits_shape[3]  # Time indices
          + target_fsas_vec.unit_positions.to(torch.int64)
          * logits_shape[3]  # Units (text) indices
          + target_fsas_vec.labels.to(torch.int64)  # Labels
      )
      return indices
  
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 7 for calling the GraphTransducerLossBase obj's underlying member method code with example usage: obj.get_logits_indices(target_fsas_vec, logits_shape) and returns: indices