code: def _add_token_emb_to_gate_inputs(
      gate_inputs, token_embeddings, moe_word_embed_mode
  ):
      """Add token_embeddings to gate_inputs based on moe_word_embed_mode."""
  
      token_embeddings = mtf.to_float(token_embeddings)
      if moe_word_embed_mode == "concat":
          gate_inputs = mtf.concat(
              [gate_inputs, token_embeddings],
              gate_inputs.shape.dims[-1].name,
          )
      elif moe_word_embed_mode == "concat_stop_grad":
          token_embeddings = mtf.stop_gradient(token_embeddings)
          gate_inputs = mtf.concat(
              [gate_inputs, token_embeddings],
              gate_inputs.shape.dims[-1].name,
          )
      elif moe_word_embed_mode == "add":
          gate_inputs += token_embeddings
      elif moe_word_embed_mode == "add_stop_grad":
          gate_inputs += mtf.stop_gradient(token_embeddings)
      elif moe_word_embed_mode == "embed_only":
          gate_inputs = token_embeddings
      else:
          raise ValueError(
              "Unimplemented moe word embed mode: {}".format(
                  moe_word_embed_mode
              )
          )
      return gate_inputs
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 10 for calling a global function with example usage: _add_token_emb_to_gate_inputs(gate_inputs, token_embeddings, moe_word_embed_mode) and returns: gate_inputs