code: def attention(
      q,
      k,
      v,
      memory_length_dim,
      key_dim,
      value_dim,
      bias=None,
      dropout_rate=0.0,
      dropout_broadcast_dims=None,
      extra_logit=None,
      context=None,
      float32_logits=True,
      z_loss_coeff=None,
  ):
      """Dot-product attention - doesn't use positional dimensions.
  
      key_dim is a Dimension representing the channels in the queries and keys
      value_dim is a Dimension representing the channels in values
      memory_length_dim is a Dimension representing the different key/value pairs.
  
      Dimensions of q: other_query_dims + {key_dim}
      Dimensions of k: other_memory_dims + {memory_length_dim, key_dim}
      Dimensions of v: other_memory_dims + {memory_length_dim, value_dim}
      other_memory_dims is a subset of other_query_dims
  
      Typically, other_query_dims={batch, heads, length}
      Typically, other_memory_dims={batch, heads}
  
      Args:
        q: a Tensor
        k: a Tensor
        v: a Tensor
        memory_length_dim: a Dimension
        key_dim: a Dimension
        value_dim: a Dimension
        bias: a Tensor to be added into the attention logits.
        dropout_rate: a float.
        dropout_broadcast_dims: an optional list of mtf.Dimension
        extra_logit: an optional scalar or tensor
        context: an optional Transformer.Context
        float32_logits: a boolean - if True, then compute logits in float32 to avoid
          numerical issues with bfloat16
        z_loss_coeff: a float, if z_loss_coeff is not None then add an auxiliary
          loss to push the attention logits closer to zero. This helps to stabilize
          model training.
  
      Returns:
        Tensor with shape q.shape - key_dim + value_dim
      """
      orig_q_shape = q.shape
      q, k, v, bias = maybe_reshape_attention_input_for_2d_sharding(
          context, q, k, v, bias, [key_dim, value_dim]
      )
      if float32_logits:
          k = mtf.cast(k, tf.float32)
          q = mtf.cast(q, tf.float32)
      logits = mtf.layers.us_einsum([q, k], reduced_dims=[key_dim])
      if bias is not None:
          logits += mtf.cast(bias, logits.dtype)
  
      # Adds auxiliary z-loss to push the attention logits towards zero.
      if z_loss_coeff is not None and context.train:
          tf.logging.info(
              "attention z_loss being added: {}".format(
                  tf.get_variable_scope().name
              )
          )
          log_z = mtf.reduce_logsumexp(logits, memory_length_dim)
          z_loss = mtf.square(log_z) * mtf.cast(
              context.nonpadding, log_z.dtype
          )
          z_loss = mtf.reduce_mean(z_loss)
          if context.num_microbatches and context.num_microbatches > 1:
              tf.logging.info(
                  "Dividing attention z-loss loss by num_microbatches={}".format(
                      context.num_microbatches
                  )
              )
              z_loss /= context.num_microbatches
          if context.train:
              mtf.scalar_summary("attention_z_loss", z_loss)
          z_loss *= z_loss_coeff
          context.losses.append(mtf.cast(z_loss, v.dtype))
  
      weights = mtf.softmax(
          logits, memory_length_dim, extra_logit=extra_logit
      )
      weights = mtf.cast(weights, v.dtype)
      weights = mtf.dropout(
          weights,
          context.train,
          1.0 - dropout_rate,
          noise_shape=weights.shape - dropout_broadcast_dims,
      )
      outputs_shape = q.shape - key_dim + value_dim
      outputs = mtf.einsum([weights, v], outputs_shape)
      outputs = mtf.reshape(outputs, orig_q_shape - key_dim + value_dim)
      return outputs
  
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 for calling a global function with example usage: attention(q, k, v, memory_length_dim, key_dim, value_dim, bias, dropout_rate, dropout_broadcast_dims, extra_logit, context, float32_logits, z_loss_coeff) and returns: outputs