def multihead_self_attention_memory_compressed(x,
                                               mask_right,
                                               compression_factor,
                                               kv_channels,
                                               heads,
                                               dropout=0.0,
                                               dropout_broadcast_dims=None,
                                               master_dtype=tf.float32,
                                               slice_dtype=tf.float32,
                                               name="multihead_attention"):
  """Memory-compressed self-attention.

  The memory is first average-pooled (strided) to make it shorter by
  a factor of compression_factor.

  Args:
    x: a mtf.Tensor with shape
      [<batch_dims>, query_length, io_channels]
    mask_right: a boolean
    compression_factor: an integer
    kv_channels: a mtf.Dimension (the size of the key and value vectors)
    heads: a mtf.Dimension (the number of heads)
    dropout: a floating point value
    dropout_broadcast_dims: an optional list of mtf.Dimension
    master_dtype: a tf.dtype
    slice_dtype: a tf.dtype
    name: an optional string.

  Returns:
    A mtf.Tensor with shape [batch, query_length, io_channels]

  Raises:
    ValueError: if the dimensions do not match.
  """
  batch_dims = x.shape.dims[:-2]
  length, io_channels = x.shape.dims[-2:]
  with tf.variable_scope(name,
                         default_name="compressed_attention",
                         values=[x]):
    wq, wk, wv, wo = multihead_attention_vars(
        x.mesh, heads, io_channels, kv_channels,
        master_dtype, slice_dtype, x.dtype)
    memory_antecedent = compress_mean(x, length, compression_factor)
    memory_antecedent = rename_length_to_memory_length(memory_antecedent)
    memory_length = memory_antecedent.shape.dims[-2]
    q = mtf.einsum(
        [x, wq],
        mtf.Shape(batch_dims + [heads, length, kv_channels]))
    k = mtf.einsum(
        [memory_antecedent, wk],
        mtf.Shape(batch_dims + [heads, memory_length, kv_channels]))
    v = mtf.einsum(
        [memory_antecedent, wv],
        mtf.Shape(batch_dims + [heads, memory_length, kv_channels]))
    if mask_right:
      query_pos = mtf.range(x.mesh, length, dtype=tf.int32)
      memory_pos = (
          mtf.range(x.mesh, memory_length, dtype=tf.int32) * compression_factor
          + (compression_factor - 1))
      mask = mtf.cast(mtf.greater(memory_pos, query_pos), x.dtype) * -1e9
    else:
      mask = None
    o = dot_product_attention(
        q, k, v, mask, dropout, dropout_broadcast_dims, extra_logit=0.0)
    return mtf.einsum(
        [o, wo], mtf.Shape(batch_dims + [length, io_channels]))
def multihead_self_attention_memory_compressed ( x , mask_right , compression_factor , kv_channels , heads , dropout = 0.0 , dropout_broadcast_dims = None , master_dtype = tf . float32 , slice_dtype = tf . float32 , name = "multihead_attention" ) : batch_dims = x . shape . dims [ : - 2 ] length , io_channels = x . shape . dims [ - 2 : ] with tf . variable_scope ( name , default_name = "compressed_attention" , values = [ x ] ) : wq , wk , wv , wo = multihead_attention_vars ( x . mesh , heads , io_channels , kv_channels , master_dtype , slice_dtype , x . dtype ) memory_antecedent = compress_mean ( x , length , compression_factor ) memory_antecedent = rename_length_to_memory_length ( memory_antecedent ) memory_length = memory_antecedent . shape . dims [ - 2 ] q = mtf . einsum ( [ x , wq ] , mtf . Shape ( batch_dims + [ heads , length , kv_channels ] ) ) k = mtf . einsum ( [ memory_antecedent , wk ] , mtf . Shape ( batch_dims + [ heads , memory_length , kv_channels ] ) ) v = mtf . einsum ( [ memory_antecedent , wv ] , mtf . Shape ( batch_dims + [ heads , memory_length , kv_channels ] ) ) if mask_right : query_pos = mtf . range ( x . mesh , length , dtype = tf . int32 ) memory_pos = ( mtf . range ( x . mesh , memory_length , dtype = tf . int32 ) * compression_factor + ( compression_factor - 1 ) ) mask = mtf . cast ( mtf . greater ( memory_pos , query_pos ) , x . dtype ) * - 1e9 else : mask = None o = dot_product_attention ( q , k , v , mask , dropout , dropout_broadcast_dims , extra_logit = 0.0 ) return mtf . einsum ( [ o , wo ] , mtf . Shape ( batch_dims + [ length , io_channels ] ) )
Memory-compressed self-attention.

  The memory is first average-pooled (strided) to make it shorter by
  a factor of compression_factor.

  Args:
    x: a mtf.Tensor with shape
      [<batch_dims>, query_length, io_channels]
    mask_right: a boolean
    compression_factor: an integer
    kv_channels: a mtf.Dimension (the size of the key and value vectors)
    heads: a mtf.Dimension (the number of heads)
    dropout: a floating point value
    dropout_broadcast_dims: an optional list of mtf.Dimension
    master_dtype: a tf.dtype
    slice_dtype: a tf.dtype
    name: an optional string.

  Returns:
    A mtf.Tensor with shape [batch, query_length, io_channels]

  Raises:
    ValueError: if the dimensions do not match.