class_docstr: STACL: Simultaneous Translation with Implicit Anticipation and
  Controllable Latency using Prefix-to-Prefix Framework
  https://www.aclweb.org/anthology/P19-1289/
class_docstr_tok: ["STACL",":","Simultaneous","Translation","with","Implicit","Anticipation","and","Controllable","Latency","using","Prefix-to-Prefix","Framework","https",":","\/\/www.aclweb.org\/anthology\/P19-1289\/"]
code: def p_choose_from_qk(
      self,
      query: Optional[Tensor],
      key: Optional[Tensor],
      key_padding_mask: Optional[Tensor] = None,
      incremental_state: Optional[
          Dict[str, Dict[str, Optional[Tensor]]]
      ] = None,
  ):
      assert query is not None
      assert key is not None
  
      p_choose = waitk_p_choose(
          tgt_len=query.size(0),
          src_len=key.size(0),
          bsz=query.size(1) * self.num_heads,
          waitk_lagging=self.waitk_lagging,
          key_padding_mask=key_padding_mask,
          incremental_state=incremental_state,
      )
  
      return p_choose.to(query)
  
imports: import math
  import torch
  from torch import Tensor
  import torch.nn
  from examples.simultaneous_translation.utils.p_choose_strategy import learnable_p_choose, waitk_p_choose
  from examples.simultaneous_translation.utils.monotonic_attention import expected_alignment_from_p_choose, expected_soft_attention, mass_preservation
  from fairseq.modules import MultiheadAttention
  from .None import register_monotonic_attention
  from typing import Dict, Optional
label_desc: Use image node_id 3 for calling the WaitKAttention obj's underlying member method code with example usage: obj.p_choose_from_qk(query, key, key_padding_mask, incremental_state) and returns: p_choose