class_docstr: Duration calculator module.
code: def _calculate_duration(att_ws):
      if len(att_ws.shape) == 2:
          # tacotron 2 case -> (T_feats, T_text)
          pass
      elif len(att_ws.shape) == 4:
          # transformer case -> (#layers, #heads, T_feats, T_text)
          # get the most diagonal head according to focus rate
          att_ws = torch.cat(
              [att_w for att_w in att_ws], dim=0
          )  # (#heads * #layers, T_feats, T_text)
          diagonal_scores = att_ws.max(dim=-1)[0].mean(
              dim=-1
          )  # (#heads * #layers,)
          diagonal_head_idx = diagonal_scores.argmax()
          att_ws = att_ws[diagonal_head_idx]  # (T_feats, T_text)
      else:
          raise ValueError(
              "att_ws should be 2 or 4 dimensional tensor."
          )
      # calculate duration from 2d attention weight
      durations = torch.stack(
          [
              att_ws.argmax(-1).eq(i).sum()
              for i in range(att_ws.shape[1])
          ]
      )
      return durations.view(-1)
  
imports: from typing import Tuple
  import torch
label_desc: Use image node_id 3 for calling the DurationCalculator obj's underlying member method code with example usage: obj._calculate_duration(att_ws) and returns: durations