class_docstr: Create attention parameters using experts-layer.
docstr: Computes qkv all in one call using MoE layer.
docstr_tok: ["Computes","qkv","all","in","one","call","using","MoE","layer","."]
code: def _compute_merge_qkv(self, antecedent):
      """Computes qkv all in one call using MoE layer."""
  
      def _replace_d_model_dim(t):
          """Used to replace the `d_model` dim with `heads`."""
          new_last_dim = mtf.Dimension(
              self.q_shape[-1].name, t.shape[-1].size
          )
          return mtf.reshape(
              t, new_shape=mtf.Shape(t.shape[:-1] + [new_last_dim])
          )
  
      if self.expert_computation == "qkv":
          # NOTE: This assumes querty and memory antecedent are the same
          qk = self.moe_layer.call(self.context, antecedent)
          # Split qk here since they went through experts-layers
          q, k = qk
          q = _replace_d_model_dim(q)
          k = _replace_d_model_dim(k)
      elif self.expert_computation == "q":
          q = self.moe_layer.call(self.context, antecedent)
          q = _replace_d_model_dim(q)
          # Compute key/value normally
          k = mtf.layers.us_einsum(
              [antecedent, self.wkv],
              reduced_dims=[self.memory_input_dim],
          )
      elif self.expert_computation == "kv":
          k = self.moe_layer.call(self.context, antecedent)
          k = _replace_d_model_dim(k)
          # Compute query normally
          q = mtf.layers.us_einsum(
              [antecedent, self.wq], reduced_dims=[self.query_input_dim]
          )
      else:
          raise ValueError(
              "Invalid expert computation mode: {}".format(
                  self.expert_computation
              )
          )
  
      # Scale query
      q *= self.key_dim.size**-0.5
      self._q = mtf.replace_dimensions(q, q.shape.dims[-1], self.q_dims)
      self._k = mtf.replace_dimensions(k, k.shape.dims[-1], self.k_dims)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling the ExpertsAttentionParams obj's underlying member method code with example usage: obj._compute_merge_qkv(antecedent) and returns: mtf