class_docstr: Multihead Attention Encoder
  
  Args:
      hparams (dict or HParams, optional): Hyperparameters. Missing
          hyperparamerter will be set to default values. See
          :meth:`default_hparams` for the hyperparameter sturcture and
          default values.
  
  .. document private functions
  .. automethod:: _build
class_docstr_tok: ["Multihead","Attention","Encoder","Args",":","hparams","(","dict","or","HParams",",","optional",")",":","Hyperparameters",".","Missing","hyperparamerter","will","be","set","to","default","values",".","See",":","meth",":","`","default_hparams","`","for","the","hyperparameter","sturcture","and","default","values",".","..","document","private","functions","..","automethod",":",":","_build"]
docstr: Encodes the inputs.
  
  Args:
      queries: A 3d tensor with shape of [batch, length_query,
          depth_query].
      memory: A 3d tensor with shape of [batch, length_key, depth_key].
      memory_attention_bias: A 3d tensor with shape of
          [batch, length_key, num_units].
      cache: Memory cache only when inferencing the sentence from sractch.
      mode (optional): A tensor taking value in
          :tf_main:`tf.estimator.ModeKeys <estimator/ModeKeys>`, including
          `TRAIN`, `EVAL` and `PREDICT`. Controls dropout mode.
          If `None` (default), :func:`texar.tf.global_mode` is used.
  
  Returns:
      A Tensor of shape `[batch_size, max_time, dim]` containing the
      encoded vectors.
docstr_tok: ["Encodes","the","inputs",".","Args",":","queries",":","A","3d","tensor","with","shape","of","[","batch",",","length_query",",","depth_query","]",".","memory",":","A","3d","tensor","with","shape","of","[","batch",",","length_key",",","depth_key","]",".","memory_attention_bias",":","A","3d","tensor","with","shape","of","[","batch",",","length_key",",","num_units","]",".","cache",":","Memory","cache","only","when","inferencing","the","sentence","from","sractch",".","mode","(","optional",")",":","A","tensor","taking","value","in",":","tf_main",":","`","tf.estimator.ModeKeys","<","estimator\/ModeKeys",">","`",",","including","`","TRAIN","`",",","`","EVAL","`","and","`","PREDICT","`",".","Controls","dropout","mode",".","If","`","None","`","(","default",")",",",":","func",":","`","texar.tf.global_mode","`","is","used",".","Returns",":","A","Tensor","of","shape","`","[","batch_size",",","max_time",",","dim","]","`","containing","the","encoded","vectors","."]
code: def _build(
      self,
      queries,
      memory,
      memory_attention_bias,
      cache=None,
      mode=None,
  ):
      """Encodes the inputs.
  
      Args:
          queries: A 3d tensor with shape of [batch, length_query,
              depth_query].
          memory: A 3d tensor with shape of [batch, length_key, depth_key].
          memory_attention_bias: A 3d tensor with shape of
              [batch, length_key, num_units].
          cache: Memory cache only when inferencing the sentence from sractch.
          mode (optional): A tensor taking value in
              :tf_main:`tf.estimator.ModeKeys <estimator/ModeKeys>`, including
              `TRAIN`, `EVAL` and `PREDICT`. Controls dropout mode.
              If `None` (default), :func:`texar.tf.global_mode` is used.
  
      Returns:
          A Tensor of shape `[batch_size, max_time, dim]` containing the
          encoded vectors.
      """
  
      with tf.variable_scope(self.variable_scope):
          num_heads = self._hparams.num_heads
          num_units = self._hparams.num_units
          if num_units % num_heads:
              raise ValueError(
                  "Value depth (%d) must be divisible by "
                  "the number of attention heads (%d)."
                  % (num_units, num_heads)
              )
  
          def _update_and_return(layer, key):
              if memory is None:
                  # Self Attention
                  out = layer(queries)
  
                  if cache is not None:
                      # 'decoder self attention when dynamic decoding'
                      key = "self_{}".format(key)
                      res = cache[key]
                      if isinstance(res, tf.TensorArray):
                          # inference-like decoding
                          # TODO(zhiting): This writing op may cause a bug
                          # on CPU--it looks the two TensorArray
                          # cache['self_keys'] and cache['self_values']
                          # will mix up starting from certain step, causing
                          # shape mismatch. This op looks fine on GPU.
                          res = res.write(
                              res.size(), tf.squeeze(out, axis=[1])
                          )
                          out = transpose_batch_time(res.stack())
                      else:
                          # normal decoding
                          res = tf.concat([res, out], axis=1)
                          out = res
                      cache[key] = res
  
              else:
                  # encoder decoder attention
                  if cache is not None:
                      key = "memory_{}".format(key)
                      res = cache[key]
                      if isinstance(res, tf.TensorArray):
                          # inference-like decoding
                          size = res.size()
                          false_fn = lambda: transpose_batch_time(
                              res.stack()
                          )
                      else:
                          # normal decoding
                          size = tf.shape(res)[1]
                          false_fn = lambda: res
                      out = tf.cond(
                          tf.equal(size, 0),
                          true_fn=lambda: layer(memory),
                          false_fn=false_fn,
                      )
                  else:
                      out = layer(memory)
  
              return out
  
          Q = self.Q_dense(queries)
          K = _update_and_return(self.K_dense, "keys")
          V = _update_and_return(self.V_dense, "values")
  
          Q_ = self._split_heads(Q)
          K_ = self._split_heads(K)
          V_ = self._split_heads(V)
          # [batch_size, num_heads, seq_length, memory_depth]
          key_depth_per_head = num_units // num_heads
          Q_ *= key_depth_per_head**-0.5
  
          logits = tf.matmul(Q_, K_, transpose_b=True)
          if memory_attention_bias is not None:
              logits += memory_attention_bias
          weights = tf.nn.softmax(logits, name="attention_weights")
          weights = tf.layers.dropout(
              weights,
              rate=self._hparams.dropout_rate,
              training=is_train_mode(mode),
          )
          outputs = tf.matmul(weights, V_)
  
          outputs = self._combine_heads(outputs)
          outputs = self.O_dense(outputs)
          # (batch_size, length_query, output_dim)
  
      if not self._built:
          self._add_internal_trainable_variables()
          self._built = True
  
      return outputs
  
file_path: texar/texar/tf/modules/encoders/multihead_attention.py
imports: import tensorflow
  from texar.tf.core import layers
  from texar.tf.modules.encoders.encoder_base import EncoderBase
  from texar.tf.utils.shapes import shape_list
  from texar.tf.utils.mode import is_train_mode
  from texar.tf.utils import transpose_batch_time
label_desc: Use image node_id 3 for calling the MultiheadAttentionEncoder obj's underlying member method code with example usage: obj._build(queries, memory, memory_attention_bias, cache, mode) and returns: outputs, out