original_docstring: 
              Compute a Scaled Dot Product Attention.
          :param query: size: batch, head, 1 (ego-entity), features
          :param key:  size: batch, head, entities, features
          :param value: size: batch, head, entities, features
          :param mask: size: batch,  head, 1 (absence feature), 1 (ego-entity)
          :return: the attention softmax(QK^T/sqrt(dk))V
          
docstring: Compute a Scaled Dot Product Attention.
docstring_tokens: Compute a Scaled Dot Product Attention .
code: def attention(self, query, key, value, mask):
          """
              Compute a Scaled Dot Product Attention.
          :param query: size: batch, head, 1 (ego-entity), features
          :param key:  size: batch, head, entities, features
          :param value: size: batch, head, entities, features
          :param mask: size: batch,  head, 1 (absence feature), 1 (ego-entity)
          :return: the attention softmax(QK^T/sqrt(dk))V
          """
          d_k = self.features_per_head
          scores = tf.matmul(query, tf.transpose(key, perm=[0, 1, 3, 2])) / np.sqrt(d_k)
          mask_constant = scores * 0 + -1e9
          if mask is not None:
              scores = tf.where(mask, mask_constant, scores)
          p_attn = tf.nn.softmax(scores, dim=-1)
          att_output = tf.matmul(p_attn, value)
          return att_output, p_attn
code_tokens: def attention ( self , query , key , value , mask ) : d_k = self . features_per_head scores = tf . matmul ( query , tf . transpose ( key , perm = [ 0 , 1 , 3 , 2 ] ) ) / np . sqrt ( d_k ) mask_constant = scores * 0 + - 1e9 if mask is not None : scores = tf . where ( mask , mask_constant , scores ) p_attn = tf . nn . softmax ( scores , dim = - 1 ) att_output = tf . matmul ( p_attn , value ) return att_output , p_attn
short_docstring: Compute a Scaled Dot Product Attention.
short_docstring_tokens: Compute a Scaled Dot Product Attention .
comment: """
              Compute a Scaled Dot Product Attention.
          :param query: size: batch, head, 1 (ego-entity), features
          :param key:  size: batch, head, entities, features
          :param value: size: batch, head, entities, features
          :param mask: size: batch,  head, 1 (absence feature), 1 (ego-entity)
          :return: the attention softmax(QK^T/sqrt(dk))V
          """