docstr: input:
      feats: (batch, seq_len, self.tag_size+2)
      masks: (batch, seq_len)
docstr_tok: ["input",":","feats",":","(","batch",",","seq_len",",","self.tag_size+2",")","masks",":","(","batch",",","seq_len",")"]
code: def _calculate_PZ(self, feats, mask):
      """
      input:
          feats: (batch, seq_len, self.tag_size+2)
          masks: (batch, seq_len)
      """
      batch_size = feats.size(0)
      seq_len = feats.size(1)
      tag_size = feats.size(2)
      # print feats.view(seq_len, tag_size)
      assert tag_size == self.tagset_size + 2
      mask = mask.transpose(1, 0).contiguous()
      ins_num = seq_len * batch_size
      ## be careful the view shape, it is .view(ins_num, 1, tag_size) but not .view(ins_num, tag_size, 1)
      feats = (
          feats.transpose(1, 0)
          .contiguous()
          .view(ins_num, 1, tag_size)
          .expand(ins_num, tag_size, tag_size)
      )
      ## need to consider start
      scores = feats + self.transitions.view(
          1, tag_size, tag_size
      ).expand(ins_num, tag_size, tag_size)
      scores = scores.view(seq_len, batch_size, tag_size, tag_size)
      # build iter
      seq_iter = enumerate(scores)
      _, inivalues = next(
          seq_iter
      )  # bat_size * from_target_size * to_target_size
      # only need start from start_tag
      partition = (
          inivalues[:, self.START_TAG, :]
          .clone()
          .view(batch_size, tag_size, 1)
      )  # bat_size * to_target_size
  
      ## add start score (from start to all tag, duplicate to batch_size)
      # partition = partition + self.transitions[START_TAG,:].view(1, tag_size, 1).expand(batch_size, tag_size, 1)
      # iter over last scores
      for idx, cur_values in seq_iter:
          # previous to_target is current from_target
          # partition: previous results log(exp(from_target)), #(batch_size * from_target)
          # cur_values: bat_size * from_target * to_target
  
          cur_values = cur_values + partition.contiguous().view(
              batch_size, tag_size, 1
          ).expand(batch_size, tag_size, tag_size)
          cur_partition = log_sum_exp(cur_values, tag_size)
          # print cur_partition.data
  
          # (bat_size * from_target * to_target) -> (bat_size * to_target)
          # partition = utils.switch(partition, cur_partition, mask[idx].view(bat_size, 1).expand(bat_size, self.tagset_size)).view(bat_size, -1)
          mask_idx = (
              mask[idx, :]
              .view(batch_size, 1)
              .expand(batch_size, tag_size)
          )
  
          ## effective updated partition part, only keep the partition value of mask value = 1
          masked_cur_partition = cur_partition.masked_select(mask_idx)
          ## let mask_idx broadcastable, to disable warning
          mask_idx = mask_idx.contiguous().view(batch_size, tag_size, 1)
  
          ## replace the partition where the maskvalue=1, other partition value keeps the same
          partition.masked_scatter_(mask_idx, masked_cur_partition)
          # until the last state, add transition score for all partition (and do log_sum_exp) then select the value in STOP_TAG
      cur_values = self.transitions.view(1, tag_size, tag_size).expand(
          batch_size, tag_size, tag_size
      ) + partition.contiguous().view(batch_size, tag_size, 1).expand(
          batch_size, tag_size, tag_size
      )
  
      cur_partition = log_sum_exp(cur_values, tag_size)
      final_partition = cur_partition[:, self.STOP_TAG]
      return final_partition.sum(), scores
  
file_path: DeepKE/src/deepke/event_extraction/standard/bertcrf/crf.py
imports: import torch
  import torch.autograd
  import torch.nn
label_desc: Use image node_id 3 for calling the CRF obj's underlying member method code with example usage: obj._calculate_PZ(feats, mask) and returns: final_partition,  scores