class_docstr: DotAttention
class_docstr_tok: ["DotAttention"]
docstr: :param s: [src_sequence_length, batch_size, src_dim]
  :param h: [batch_size, tgt_dim] or [tgt_sequence_length, batch_size, tgt_dim]
  :param mask: [src_sequence_length, batch_size]             or [tgt_sequence_length, src_sequence_length, batch_sizse]
  :param pre_compute: [src_sequence_length, batch_size, hidden_dim]
  :return: [src_sequence_length, batch_size]             or [tgt_sequence_length, src_sequence_length, batch_size]
docstr_tok: [":","param","s",":","[","src_sequence_length",",","batch_size",",","src_dim","]",":","param","h",":","[","batch_size",",","tgt_dim","]","or","[","tgt_sequence_length",",","batch_size",",","tgt_dim","]",":","param","mask",":","[","src_sequence_length",",","batch_size","]","or","[","tgt_sequence_length",",","src_sequence_length",",","batch_sizse","]",":","param","pre_compute",":","[","src_sequence_length",",","batch_size",",","hidden_dim","]",":","return",":","[","src_sequence_length",",","batch_size","]","or","[","tgt_sequence_length",",","src_sequence_length",",","batch_size","]"]
code: def get_prob(self, src, tgt, mask, pre_compute, return_logits=False):
      """
          :param s: [src_sequence_length, batch_size, src_dim]
          :param h: [batch_size, tgt_dim] or [tgt_sequence_length, batch_size, tgt_dim]
          :param mask: [src_sequence_length, batch_size]\
               or [tgt_sequence_length, src_sequence_length, batch_sizse]
          :param pre_compute: [src_sequence_length, batch_size, hidden_dim]
          :return: [src_sequence_length, batch_size]\
               or [tgt_sequence_length, src_sequence_length, batch_size]
          """
      s_shape = src.get_shape().as_list()
      h_shape = tgt.get_shape().as_list()
      src_dim = s_shape[-1]
      tgt_dim = h_shape[-1]
      assert src_dim is not None, "src dimension must be defined"
      assert tgt_dim is not None, "tgt dimension must be defined"
  
      self._define_params(src_dim, tgt_dim)
  
      if len(h_shape) == 2:
          tgt = tf.expand_dims(tgt, 0)
      if pre_compute is None:
          pre_compute = self.get_pre_compute(src)
  
      buf0 = pre_compute
      buf1 = tf.tensordot(tgt, self.var["U"], axes=[[2], [0]])
      buf2 = tf.tanh(tf.expand_dims(buf0, 0) + tf.expand_dims(buf1, 1))
  
      if not self.is_vanilla:
          xh1 = tgt
          xh2 = tgt
          s1 = src
          if self.need_padding:
              xh1 = tf.tensordot(xh1, self.var["V_t"], 1)
              xh2 = tf.tensordot(xh2, self.var["S_t"], 1)
              s1 = tf.tensordot(s1, self.var["V_s"], 1)
          if not self.is_identity_transform:
              xh1 = tf.tensordot(xh1, self.var["T"], 1)
              xh2 = tf.tensordot(xh2, self.var["T"], 1)
          buf3 = tf.expand_dims(s1, 0) * tf.expand_dims(xh1, 1)
          buf3 = tf.tanh(
              tf.tensordot(buf3, self.var["V"], axes=[[3], [0]])
          )
          buf = tf.reshape(tf.tanh(buf2 + buf3), shape=tf.shape(buf3))
      else:
          buf = buf2
      v = self.var["v"]
      e = tf.tensordot(buf, v, [[3], [0]])
      e = tf.squeeze(e, axis=[3])
      tmp = tf.reshape(e + (mask - 1) * 10000.0, shape=tf.shape(e))
      prob = tf.nn.softmax(tmp, 1)
      if len(h_shape) == 2:
          prob = tf.squeeze(prob, axis=[0])
          tmp = tf.squeeze(tmp, axis=[0])
      if return_logits:
          return prob, tmp
      return prob
  
file_path: nni/examples/trials/weight_sharing/ga_squad/attention.py
imports: import math
  import tensorflow
  from tensorflow.python.ops.rnn_cell_impl import RNNCell
label_desc: Use image node_id 11 for calling the DotAttention obj's underlying member method code with example usage: obj.get_prob(src, tgt, mask, pre_compute, return_logits) and returns: prob, prob,  tmp