class_docstr: Conditional random field.
  This module implements a conditional random field [LMP01]_. The forward computation
  of this class computes the log likelihood of the given sequence of tags and
  emission score tensor. This class also has `~CRF.decode` method which finds
  the best tag sequence given an emission score tensor using `Viterbi algorithm`_.
  Args:
      num_tags: Number of tags.
      batch_first: Whether the first dimension corresponds to the size of a minibatch.
  Attributes:
      start_transitions (`~torch.nn.Parameter`): Start transition score tensor of size
          ``(num_tags,)``.
      end_transitions (`~torch.nn.Parameter`): End transition score tensor of size
          ``(num_tags,)``.
      transitions (`~torch.nn.Parameter`): Transition score tensor of size
          ``(num_tags, num_tags)``.
  .. [LMP01] Lafferty, J., McCallum, A., Pereira, F. (2001).
     "Conditional random fields: Probabilistic models for segmenting and
     labeling sequence data". *Proc. 18th International Conf. on Machine
     Learning*. Morgan Kaufmann. pp. 282–289.
  .. _Viterbi algorithm: https://en.wikipedia.org/wiki/Viterbi_algorithm
class_docstr_tok: ["Conditional","random","field",".","This","module","implements","a","conditional","random","field","[","LMP01","]","_",".","The","forward","computation","of","this","class","computes","the","log","likelihood","of","the","given","sequence","of","tags","and","emission","score","tensor",".","This","class","also","has","`","~CRF.decode","`","method","which","finds","the","best","tag","sequence","given","an","emission","score","tensor","using","`","Viterbi","algorithm","`","_",".","Args",":","num_tags",":","Number","of","tags",".","batch_first",":","Whether","the","first","dimension","corresponds","to","the","size","of","a","minibatch",".","Attributes",":","start_transitions","(","`","~torch.nn.Parameter","`",")",":","Start","transition","score","tensor","of","size","``","(","num_tags",",",")","``",".","end_transitions","(","`","~torch.nn.Parameter","`",")",":","End","transition","score","tensor","of","size","``","(","num_tags",",",")","``",".","transitions","(","`","~torch.nn.Parameter","`",")",":","Transition","score","tensor","of","size","``","(","num_tags",",","num_tags",")","``",".","..","[","LMP01","]","Lafferty",",","J.",",","McCallum",",","A.",",","Pereira",",","F.","(","2001",")",".","``","Conditional","random","fields",":","Probabilistic","models","for","segmenting","and","labeling","sequence","data","''",".","*","Proc",".","18th","International","Conf",".","on","Machine","Learning","*",".","Morgan","Kaufmann",".","pp",".","282\u2013289",".","..","_Viterbi","algorithm",":","https",":","\/\/en.wikipedia.org\/wiki\/Viterbi_algorithm"]
code: def _viterbi_decode(
      self,
      emissions: torch.FloatTensor,
      mask: torch.ByteTensor,
      pad_tag: Optional[int] = None,
  ) -> List[List[int]]:
      # emissions: (seq_length, batch_size, num_tags)
      # mask: (seq_length, batch_size)
      # return: (batch_size, seq_length)
      if pad_tag is None:
          pad_tag = 0
  
      device = emissions.device
      seq_length, batch_size = mask.shape
  
      # Start transition and first emission
      # shape: (batch_size, num_tags)
      score = self.start_transitions + emissions[0]
      history_idx = torch.zeros(
          (seq_length, batch_size, self.num_tags),
          dtype=torch.long,
          device=device,
      )
      oor_idx = torch.zeros(
          (batch_size, self.num_tags), dtype=torch.long, device=device
      )
      oor_tag = torch.full(
          (seq_length, batch_size),
          pad_tag,
          dtype=torch.long,
          device=device,
      )
  
      # - score is a tensor of size (batch_size, num_tags) where for every batch,
      #   value at column j stores the score of the best tag sequence so far that ends
      #   with tag j
      # - history_idx saves where the best tags candidate transitioned from; this is used
      #   when we trace back the best tag sequence
      # - oor_idx saves the best tags candidate transitioned from at the positions
      #   where mask is 0, i.e. out of range (oor)
  
      # Viterbi algorithm recursive case: we compute the score of the best tag sequence
      # for every possible next tag
      # print(f'0 step score: {score[0, -3:]}') #debug
      for i in range(1, seq_length):
          # Broadcast viterbi score for every possible next tag
          # shape: (batch_size, num_tags, 1)
          broadcast_score = score.unsqueeze(2)
  
          # Broadcast emission score for every possible current tag
          # shape: (batch_size, 1, num_tags)
          broadcast_emission = emissions[i].unsqueeze(1)
  
          # Compute the score tensor of size (batch_size, num_tags, num_tags) where
          # for each sample, entry at row i and column j stores the score of the best
          # tag sequence so far that ends with transitioning from tag i to tag j and emitting
          # shape: (batch_size, num_tags, num_tags)
          next_score = (
              broadcast_score + self.transitions + broadcast_emission
          )
  
          # Find the maximum score over all possible current tag
          # shape: (batch_size, num_tags)
          next_score, indices = next_score.max(dim=1)
  
          # Set score to the next score if this timestep is valid (mask == 1)
          # and save the index that produces the next score
          # shape: (batch_size, num_tags)
          score = torch.where(mask[i].unsqueeze(-1), next_score, score)
          indices = torch.where(mask[i].unsqueeze(-1), indices, oor_idx)
          history_idx[i - 1] = indices
  
      # End transition score
      # shape: (batch_size, num_tags)
      end_score = score + self.end_transitions
      best_score, end_tag = end_score.max(dim=1)
      """
          print(f'history_idx: {history_idx[0, :, :]}')
          print(f"pt end score: {end_score[0, :]}")  # debug
          print(
              f"pt best score: {end_score[0, end_tag[0]]}, last tag: {end_tag[0]}"
          )  # debug
          """
  
      # shape: (batch_size,)
      seq_ends = mask.long().sum(dim=0) - 1
  
      # insert the best tag at each sequence end (last position with mask == 1)
      history_idx = history_idx.transpose(1, 0).contiguous()
      history_idx.scatter_(
          1,
          seq_ends.view(-1, 1, 1).expand(-1, 1, self.num_tags),
          end_tag.view(-1, 1, 1).expand(-1, 1, self.num_tags),
      )
      history_idx = history_idx.transpose(1, 0).contiguous()
  
      # The most probable path for each sequence
      best_tags_arr = torch.zeros(
          (seq_length, batch_size), dtype=torch.long, device=device
      )
      best_tags = torch.zeros(
          batch_size, 1, dtype=torch.long, device=device
      )
      for idx in range(seq_length - 1, -1, -1):
          best_tags = torch.gather(history_idx[idx], 1, best_tags)
          best_tags_arr[idx] = best_tags.data.view(batch_size)
  
      return (
          torch.where(mask, best_tags_arr, oor_tag).transpose(0, 1),
          best_score,
      )
  
file_path: lightseq/tests/torch_crf.py
imports: import torch
  import torch.nn
  from typing import List, Optional
label_desc: Use image node_id 9 for calling the CRF obj's underlying member method code with example usage: obj._viterbi_decode(emissions, mask, pad_tag) and returns: torch,  best_score