code: def attention_params_simple(
      mesh, io_dim, kv_dim, heads_dim, variable_dtype
  ):
      """Common case attention parameters.
  
      Args:
        mesh: a Mesh
        io_dim: a Dimension (channels dimension of inputs and outputs)
        kv_dim: a Dimension (channels in keys and values)
        heads_dim: a Dimension (number of attention "heads")
        variable_dtype: a mtf.VariableDType
      Returns:
        an AttentionParams
      """
      return AttentionParams(
          mesh,
          query_input_dim=io_dim,
          memory_input_dim=io_dim,
          output_dim=io_dim,
          key_dim=kv_dim,
          value_dim=kv_dim,
          query_heads_dims=[heads_dim],
          memory_heads_dims=[heads_dim],
          variable_dtype=variable_dtype,
      )
  
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 5 for calling a global function with example usage: attention_params_simple(mesh, io_dim, kv_dim, heads_dim, variable_dtype) and returns: AttentionParams