class_docstr: Create attention parameters using experts-layer.
code: def __init__(
      self,
      mesh,
      query_input_dim,
      memory_input_dim,
      output_dim,
      key_dim,
      value_dim,
      query_heads_dims,
      memory_heads_dims,
      variable_dtype,
      shared_kv=False,
      no_query=False,
      combine_dims=True,
      ensemble_dim=None,
      keep_query_heads_dims=False,
      fold_scaling_into_initializer=True,
      context=None,
      experts_hparams=None,
      expert_computation="qkv",
  ):
      super(ExpertsAttentionParams, self).__init__(
          mesh=mesh,
          query_input_dim=query_input_dim,
          memory_input_dim=memory_input_dim,
          output_dim=output_dim,
          key_dim=key_dim,
          value_dim=value_dim,
          query_heads_dims=query_heads_dims,
          memory_heads_dims=memory_heads_dims,
          variable_dtype=variable_dtype,
          shared_kv=shared_kv,
          no_query=no_query,
          combine_dims=combine_dims,
          ensemble_dim=ensemble_dim,
          keep_query_heads_dims=keep_query_heads_dims,
          fold_scaling_into_initializer=fold_scaling_into_initializer,
          make_attention_vars=False,
      )
  
      self.context = context
      self.expert_computation = expert_computation
  
      # Unless we want to compute both q and kv, we can use the normal MoE
      # settings.
      if expert_computation == "qkv":
          experts_attention_compute_qkv = True
      elif expert_computation in ["q", "kv"]:
          experts_attention_compute_qkv = False
          if expert_computation == "q":
              # Always assume shared_kv.
              self.wkv = mtf.get_variable(
                  self.mesh,
                  "kv",
                  self.k_shape,
                  initializer=tf.random_normal_initializer(
                      stddev=self.memory_input_dim.size**-0.5
                  ),
                  dtype=self.variable_dtype,
              )
          else:  # Computing kv with experts.
              self.wq = mtf.get_variable(
                  self.mesh,
                  "q",
                  self.q_shape,
                  initializer=tf.random_normal_initializer(
                      stddev=self.query_input_dim.size**-0.5
                  ),
                  dtype=self.variable_dtype,
              )
      else:
          raise ValueError(
              "Invalid expert computation mode: {}".format(
                  expert_computation
              )
          )
  
      # ExpertsAttention, for simplicitly, asserts that combine_dims is True, and
      # for efficiency, that shared_kv is True.
      if not self.combine_dims:
          raise ValueError(
              "combine_dims must be True for ExpertsAttention."
          )
      if not self.shared_kv:
          raise ValueError(
              "shared_kv must be True for ExpertsAttention."
          )
      if mtf.layers.unit_scaling_convention():
          raise NotImplementedError
  
      # Now replace "heads" dim with the "d_model" name to avoid conflicts when
      # we want to partition both "experts_hidden" and "heads".
      moe_output_dims = mtf.Dimension("d_model", self.q_shape[-1].size)
  
      tf.logging.info(
          "ExpertsAttention moe_hidden_size: {}".format(
              experts_hparams.hidden_size
          )
      )
      tf.logging.info("moe_output_dims: {}".format(moe_output_dims))
      self.moe_layer = mtf.transformer.moe.MoE1D(
          moe_gating=experts_hparams.moe_gating,
          num_experts=experts_hparams.num_experts,
          loss_coef=experts_hparams.loss_coef,
          group_size=experts_hparams.group_size,
          min_expert_capacity=experts_hparams.min_expert_capacity,
          capacity_factor_train=experts_hparams.capacity_factor_train,
          capacity_factor_eval=experts_hparams.capacity_factor_eval,
          switch_policy_train=experts_hparams.switch_policy_train,
          switch_policy_eval=experts_hparams.switch_policy_eval,
          switch_dropout=experts_hparams.switch_dropout,
          switch_temperature=experts_hparams.switch_temperature,
          switch_jitter=experts_hparams.switch_jitter,
          ntlb_top_k=experts_hparams.ntlb_top_k,
          hidden_size=experts_hparams.hidden_size,
          output_dim=moe_output_dims,
          use_experts_attention=experts_attention_compute_qkv,
          activation=experts_hparams.activation,
          z_loss=experts_hparams.z_loss,
      )
  
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 1 to create a new ExpertsAttentionParams object from inherited base classes: AttentionParams with example: obj = ExpertsAttentionParams(mesh, query_input_dim, memory_input_dim, output_dim, key_dim, value_dim, query_heads_dims, memory_heads_dims, variable_dtype, shared_kv, no_query, combine_dims, ensemble_dim, keep_query_heads_dims, fold_scaling_into_initializer, context, experts_hparams, expert_computation)