code: def make_params_mtlprompt(
      task_num,
      num_heads,
      prefix_hidden_dim,
      kv_dim,
      context,
      mtlprompt_share,
      dropout_rate,
      prompt_length=None,
  ):
      """Returns the parameters for MTL-Prompt mode."""
  
      tf.logging.info("MTL-Prompt mode is ON!")
      # Get task ids for the batch from the context cache.
      task_id = context.cache["task-id"]
  
      # Remove length dim from shape [batch_size, length] -> [batch_size].
      task_id = mtf.reshape(
          task_id, task_id.shape - task_id.shape.dims[-1]
      )
      batch_size_dim = task_id.shape.get_dim_by_name("batch")
  
      task_num_dim = mtf.Dimension("task_num", task_num)
      prompt_length_dim = mtf.Dimension("memory_length", prompt_length)
      heads_dim = mtf.Dimension("heads", num_heads)
      prefix_hidden_dim = mtf.Dimension(
          "prefix_hidden", prefix_hidden_dim
      )
  
      with tf.variable_scope("prompt"):
          # Initialize projection network matrices.
  
          # MTL-Prompt-Share mode.
          hidden_shape = [context.model.model_dim, prefix_hidden_dim]
          scratch_shape = [prefix_hidden_dim, heads_dim, kv_dim]
          if not mtlprompt_share:
              # MTL-Prompt-Sep mode.
              hidden_shape = [task_num_dim] + hidden_shape
              scratch_shape = [task_num_dim] + scratch_shape
  
          w_k_up = mtf.get_variable(
              context.mesh,
              "w_k_up",
              mtf.Shape(hidden_shape),
              dtype=context.variable_dtype,
          )
          w_k_down = mtf.get_variable(
              context.mesh,
              "w_k_down",
              mtf.Shape(scratch_shape),
              dtype=context.variable_dtype,
          )
          w_v_up = mtf.get_variable(
              context.mesh,
              "w_v_up",
              mtf.Shape(hidden_shape),
              dtype=context.variable_dtype,
          )
          w_v_down = mtf.get_variable(
              context.mesh,
              "w_v_down",
              mtf.Shape(scratch_shape),
              dtype=context.variable_dtype,
          )
  
      scope_name = tf.get_variable_scope().name
      scope = scope_name.split("/")[0]
      prompts = context.shared_params[scope]["prompts"]
      prompts = mtf.layers.layer_norm(
          prompts, dim=context.model.model_dim, name="prompts_layernorm"
      )
      if context.train and dropout_rate != 0.0:
          prompts = mtf.dropout(
              prompts, context.train, 1.0 - dropout_rate
          )
  
      prompts_k_hidden = mtf.matmul(
          prompts,
          w_k_up,
          output_shape=[
              task_num_dim,
              prompt_length_dim,
              prefix_hidden_dim,
          ],
          reduced_dims=[context.model.model_dim],
      )
      prompts_k_hidden = mtf.relu(prompts_k_hidden)
      if context.train and dropout_rate != 0.0:
          prompts_k_hidden = mtf.dropout(
              prompts_k_hidden,
              context.train,
              keep_prob=1.0 - dropout_rate,
          )
      prompts_k = mtf.matmul(
          prompts_k_hidden,
          w_k_down,
          output_shape=[
              task_num_dim,
              prompt_length_dim,
              heads_dim,
              kv_dim,
          ],
          reduced_dims=[prefix_hidden_dim],
      )
  
      prompts_v_hidden = mtf.matmul(
          prompts,
          w_v_up,
          output_shape=[
              task_num_dim,
              prompt_length_dim,
              prefix_hidden_dim,
          ],
          reduced_dims=[context.model.model_dim],
      )
      prompts_v_hidden = mtf.relu(prompts_v_hidden)
      if context.train and dropout_rate != 0.0:
          prompts_v_hidden = mtf.dropout(
              prompts_v_hidden,
              context.train,
              keep_prob=1.0 - dropout_rate,
          )
      prompts_v = mtf.matmul(
          prompts_v_hidden,
          w_v_down,
          output_shape=[
              task_num_dim,
              prompt_length_dim,
              heads_dim,
              kv_dim,
          ],
          reduced_dims=[prefix_hidden_dim],
      )
  
      prompts_batch_k = mtf.gather(
          prompts_k,
          task_id,
          task_num_dim,
          output_shape=[
              batch_size_dim,
              prompt_length_dim,
              heads_dim,
              kv_dim,
          ],
      )
      prompts_batch_v = mtf.gather(
          prompts_v,
          task_id,
          task_num_dim,
          output_shape=[
              batch_size_dim,
              prompt_length_dim,
              heads_dim,
              kv_dim,
          ],
      )
      return prompts_batch_k, prompts_batch_v
  
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 9 for calling a global function with example usage: make_params_mtlprompt(task_num, num_heads, prefix_hidden_dim, kv_dim, context, mtlprompt_share, dropout_rate, prompt_length) and returns: prompts_batch_k,  prompts_batch_v