def masked_local_attention_1d(x,
                              kv_channels,
                              heads,
                              window_size=128,
                              master_dtype=tf.float32,
                              slice_dtype=tf.float32,
                              length_per_split=None,
                              return_kv=None,
                              params=None,
                              name=None):
  """Attention to the source position and a neighborhood to the left of it.

  Attention for a given query position p can only see memory positions
  in the range (p - window_size, p].

  Args:
    x: a mtf.Tensor with shape batch_dims + [length, io_channels]
    kv_channels: a mtf.Dimension (the size of the key and value vectors)
    heads: a mtf.Dimension (the number of heads)
    window_size: an integer
    master_dtype: a tf.dtype (deprecated - use params arg)
    slice_dtype: a tf.dtype (deprecated - use params arg)
    length_per_split: an optional integer indicating the part of the length
      dimension per processor.  You can omit if the length dimension is not
      split.
    return_kv: an optional list onto which to append the computed k and v.
    params: an optional quadruple of Tensors (see multihead_attention_params())
    name: an optional string.

  Returns:
    a Tensor with the same shape as x

  Raises:
    ValueError: if channels or depth don't match.
  """
  with tf.variable_scope(
      name, default_name="masked_local_attention_1d", values=[x]):

    batch_dims = x.shape.dims[:-2]
    length, io_channels = x.shape.dims[-2:]
    if params is None:
      wq, wk, wv, wo = multihead_attention_vars(
          x.mesh, heads, io_channels, kv_channels,
          master_dtype, slice_dtype, x.dtype)
    else:
      wq, wk, wv, wo = params

    # Get query q, keys k and values v.
    qkv_shape = mtf.Shape(batch_dims + [heads, length, kv_channels])
    q = mtf.einsum([x, wq], qkv_shape)
    k = mtf.einsum([x, wk], qkv_shape)
    v = mtf.einsum([x, wv], qkv_shape)
    if return_kv is not None:
      return_kv.extend([k, v])

    # Choose a suitable block size.
    # We choose the greatest divisor of length_per_split less than or equal
    # to max(window_size, 128)
    if length_per_split is None:
      length_per_split = length.size
    block_length = max(window_size, 128)
    while length_per_split % block_length != 0:
      block_length -= 1

    query_block_length = mtf.Dimension("query_block_length", block_length)
    memory_block_length = mtf.Dimension("memory_block_length", block_length)
    # The num_blocks dimension gets the same name as the length dimension,
    # so it will be split in the same way.
    num_blocks = mtf.Dimension(length.name, length.size // block_length)
    q_shape = batch_dims + [heads, num_blocks, query_block_length, kv_channels]
    kv_shape = batch_dims + [
        heads, num_blocks, memory_block_length, kv_channels]
    q = mtf.reshape(q, q_shape)
    k = mtf.reshape(k, kv_shape)
    v = mtf.reshape(v, kv_shape)
    # augment the keys and values for each block with keys and values for
    # the previous window_size timesteps.
    k = mtf.left_halo_exchange(k, num_blocks, memory_block_length, window_size)
    v = mtf.left_halo_exchange(v, num_blocks, memory_block_length, window_size)
    padded_memory_block_length = mtf.Dimension(
        "memory_block_length", window_size + block_length)
    mpos = mtf.range(x.mesh, padded_memory_block_length, tf.float32)
    qpos = mtf.range(x.mesh, query_block_length, tf.float32) + window_size
    # prevent looking forward
    mask = mtf.cast(mtf.greater(mpos, qpos), x.dtype) * -1e9
    # prevent looking >=block_length timesteps backward
    mask += mtf.cast(mtf.less_equal(mpos, qpos - block_length), x.dtype) * -1e9
    # Note: The first window_size-1 positions can see back into pre-time
    # where all the keys and values are zero.  We could mask this out, but we
    # don't.
    o = dot_product_attention(q, k, v, mask=mask)
    o = mtf.reshape(o, batch_dims + [heads, length, kv_channels])
    return mtf.einsum([o, wo], mtf.Shape(batch_dims + [length, io_channels]))
def masked_local_attention_1d ( x , kv_channels , heads , window_size = 128 , master_dtype = tf . float32 , slice_dtype = tf . float32 , length_per_split = None , return_kv = None , params = None , name = None ) : with tf . variable_scope ( name , default_name = "masked_local_attention_1d" , values = [ x ] ) : batch_dims = x . shape . dims [ : - 2 ] length , io_channels = x . shape . dims [ - 2 : ] if params is None : wq , wk , wv , wo = multihead_attention_vars ( x . mesh , heads , io_channels , kv_channels , master_dtype , slice_dtype , x . dtype ) else : wq , wk , wv , wo = params qkv_shape = mtf . Shape ( batch_dims + [ heads , length , kv_channels ] ) q = mtf . einsum ( [ x , wq ] , qkv_shape ) k = mtf . einsum ( [ x , wk ] , qkv_shape ) v = mtf . einsum ( [ x , wv ] , qkv_shape ) if return_kv is not None : return_kv . extend ( [ k , v ] ) if length_per_split is None : length_per_split = length . size block_length = max ( window_size , 128 ) while length_per_split % block_length != 0 : block_length -= 1 query_block_length = mtf . Dimension ( "query_block_length" , block_length ) memory_block_length = mtf . Dimension ( "memory_block_length" , block_length ) num_blocks = mtf . Dimension ( length . name , length . size // block_length ) q_shape = batch_dims + [ heads , num_blocks , query_block_length , kv_channels ] kv_shape = batch_dims + [ heads , num_blocks , memory_block_length , kv_channels ] q = mtf . reshape ( q , q_shape ) k = mtf . reshape ( k , kv_shape ) v = mtf . reshape ( v , kv_shape ) k = mtf . left_halo_exchange ( k , num_blocks , memory_block_length , window_size ) v = mtf . left_halo_exchange ( v , num_blocks , memory_block_length , window_size ) padded_memory_block_length = mtf . Dimension ( "memory_block_length" , window_size + block_length ) mpos = mtf . range ( x . mesh , padded_memory_block_length , tf . float32 ) qpos = mtf . range ( x . mesh , query_block_length , tf . float32 ) + window_size mask = mtf . cast ( mtf . greater ( mpos , qpos ) , x . dtype ) * - 1e9 mask += mtf . cast ( mtf . less_equal ( mpos , qpos - block_length ) , x . dtype ) * - 1e9 o = dot_product_attention ( q , k , v , mask = mask ) o = mtf . reshape ( o , batch_dims + [ heads , length , kv_channels ] ) return mtf . einsum ( [ o , wo ] , mtf . Shape ( batch_dims + [ length , io_channels ] ) )
Attention to the source position and a neighborhood to the left of it.

  Attention for a given query position p can only see memory positions
  in the range (p - window_size, p].

  Args:
    x: a mtf.Tensor with shape batch_dims + [length, io_channels]
    kv_channels: a mtf.Dimension (the size of the key and value vectors)
    heads: a mtf.Dimension (the number of heads)
    window_size: an integer
    master_dtype: a tf.dtype (deprecated - use params arg)
    slice_dtype: a tf.dtype (deprecated - use params arg)
    length_per_split: an optional integer indicating the part of the length
      dimension per processor.  You can omit if the length dimension is not
      split.
    return_kv: an optional list onto which to append the computed k and v.
    params: an optional quadruple of Tensors (see multihead_attention_params())
    name: an optional string.

  Returns:
    a Tensor with the same shape as x

  Raises:
    ValueError: if channels or depth don't match.