def local_attention_layer(from_tensor: tf.Tensor,
                          to_tensor: tf.Tensor,
                          from_seq_length: int,
                          to_seq_length: int,
                          attention_mask: tf.Tensor,
                          num_attention_heads: int = 1,
                          size_per_head: int = 512,
                          query_act: Optional[Text] = None,
                          key_act: Optional[Text] = None,
                          value_act: Optional[Text] = None,
                          attention_probs_dropout_prob: float = 0.0,
                          initializer_range: float = 0.02,
                          batch_size: Optional[int] = None,
                          always_attend_to_first_position: bool = True,
                          first_position_attends_to_all: bool = True,
                          attend_from_chunk_width: int = 128,
                          attend_from_chunk_stride: int = 128,
                          attend_to_chunk_width: int = 128,
                          attend_to_chunk_stride: int = 128):
  """A fork of BERT's `attention_layer` that performs local attention.

  This attention is local in that attention happens only within each block
  (as defined by the length of the stides).

  Function parameters specific to local attention (i.e. added from BERT's
  `attention_layer`) are at the bottom of the argument list.

  Args:
    from_tensor: float Tensor of shape [batch_size, from_seq_length,
      from_width].
    to_tensor: float Tensor of shape [batch_size, to_seq_length, to_width].
    from_seq_length: Length of the padded 'from' sequence in this model.
    to_seq_length: Length of the padded 'to' sequence in this model.
    attention_mask: int32 Tensor of shape [batch_size,
      from_seq_length, to_seq_length]. The values should be 1 or 0. The
      attention scores will effectively be set to -infinity for any positions in
      the mask that are 0, and will be unchanged for positions that are 1.
    num_attention_heads: int. Number of attention heads.
    size_per_head: int. Size of each attention head.
    query_act: (optional) Activation function for the query transform.
    key_act: (optional) Activation function for the key transform.
    value_act: (optional) Activation function for the value transform.
    attention_probs_dropout_prob: (optional) float. Dropout probability of the
      attention probabilities.
    initializer_range: float. Range of the weight initializer.
    batch_size: (Optional) int. If the input is 2D, this might be the batch size
      of the 3D version of the `from_tensor` and `to_tensor`.
    always_attend_to_first_position: Should all blocks be able to attend to the
      `to_tensor`'s first position (e.g. a [CLS] position)?
    first_position_attends_to_all: Should the `from_tensor`'s first position
      be able to attend to all positions within the `from_tensor`?
    attend_from_chunk_width: The width of each block-wise chunk in
      `from_tensor`.
    attend_from_chunk_stride: The number of elements to skip when moving to the
      next block in `from_tensor`.
    attend_to_chunk_width: The width of each block-wise chunk in `to_tensor`.
    attend_to_chunk_stride: The number of elements to skip when moving to the
      next block in `to_tensor`.

  Returns:
    float Tensor of shape [batch_size, from_seq_length, num_attention_heads,
      size_per_head].

  Raises:
    ValueError: Any of the arguments or tensor shapes are invalid.
  """

  if attend_from_chunk_width < attend_from_chunk_stride:
    raise ValueError("`attend_from_chunk_width` < `attend_from_chunk_stride`"
                     "would cause sequence positions to get skipped.")
  if attend_to_chunk_width < attend_to_chunk_stride:
    raise ValueError("`attend_to_chunk_width` < `attend_to_chunk_stride`"
                     "would cause sequence positions to get skipped.")

  # Create chunks that we will attend *from* and then concatenate them.
  from_chunks = []
  if first_position_attends_to_all:
    from_chunks.append((0, 1))
    # We must skip this first position so that our output sequence is the
    # correct length (this matters in the *from* sequence only).
    from_start = 1
  else:
    from_start = 0
  for chunk_start in range(from_start, from_seq_length,
                           attend_from_chunk_stride):
    chunk_end = min(from_seq_length, chunk_start + attend_from_chunk_width)
    from_chunks.append((chunk_start, chunk_end))

  # Determine the chunks that will will attend *to*.
  to_chunks = []
  if first_position_attends_to_all:
    to_chunks.append((0, to_seq_length))
  for chunk_start in range(0, to_seq_length, attend_to_chunk_stride):
    chunk_end = min(to_seq_length, chunk_start + attend_to_chunk_width)
    to_chunks.append((chunk_start, chunk_end))

  if len(from_chunks) != len(to_chunks):
    raise ValueError(
        f"Expected to have same number of `from_chunks` ({from_chunks}) and "
        f"`to_chunks` ({from_chunks}). Check strides.")

  # TODO(jhclark): Can we save a bit of extra compute by slicing the Q/K/V
  # projected versions of these instead of recomputing those projections?
  # This only helps when the Q stride isn't the same as the K/V stride.
  # Length of `attention_output_chunks` and therefore `attention_output` is
  # determined by `from_chunks` to ensure correctness. The correspondence with
  # `to_chunks` is somewhat best effort. We need to do more to enforce this.
  attention_output_chunks = []
  for (from_start, from_end), (to_start, to_end) in zip(from_chunks, to_chunks):
    from_tensor_chunk = from_tensor[:, from_start:from_end, :]
    to_tensor_chunk = to_tensor[:, to_start:to_end, :]
    # `attention_mask`: <float>[batch_size, from_seq, to_seq]
    # `attention_mask_chunk`: <float>[batch_size, from_seq_chunk, to_seq_chunk]
    attention_mask_chunk = (
        attention_mask[:, from_start:from_end, to_start:to_end])
    if always_attend_to_first_position:
      cls_attention_mask = attention_mask[:, from_start:from_end, 0:1]
      attention_mask_chunk = tf.concat(
          [cls_attention_mask, attention_mask_chunk], axis=2)

      cls_position = to_tensor[:, 0:1, :]
      to_tensor_chunk = tf.concat([cls_position, to_tensor_chunk], axis=1)
    attention_output_chunk = bert_modeling.attention_layer(
        from_tensor=from_tensor_chunk,
        to_tensor=to_tensor_chunk,
        attention_mask=attention_mask_chunk,
        num_attention_heads=num_attention_heads,
        size_per_head=size_per_head,
        query_act=query_act,
        key_act=key_act,
        value_act=value_act,
        attention_probs_dropout_prob=attention_probs_dropout_prob,
        initializer_range=initializer_range,
        batch_size=batch_size)
    attention_output_chunks.append(attention_output_chunk)
  return tf.concat(attention_output_chunks, axis=1)
A fork of BERT's `attention_layer` that performs local attention.

  This attention is local in that attention happens only within each block
  (as defined by the length of the stides).

  Function parameters specific to local attention (i.e. added from BERT's
  `attention_layer`) are at the bottom of the argument list.

  Args:
    from_tensor: float Tensor of shape [batch_size, from_seq_length,
      from_width].
    to_tensor: float Tensor of shape [batch_size, to_seq_length, to_width].
    from_seq_length: Length of the padded 'from' sequence in this model.
    to_seq_length: Length of the padded 'to' sequence in this model.
    attention_mask: int32 Tensor of shape [batch_size,
      from_seq_length, to_seq_length]. The values should be 1 or 0. The
      attention scores will effectively be set to -infinity for any positions in
      the mask that are 0, and will be unchanged for positions that are 1.
    num_attention_heads: int. Number of attention heads.
    size_per_head: int. Size of each attention head.
    query_act: (optional) Activation function for the query transform.
    key_act: (optional) Activation function for the key transform.
    value_act: (optional) Activation function for the value transform.
    attention_probs_dropout_prob: (optional) float. Dropout probability of the
      attention probabilities.
    initializer_range: float. Range of the weight initializer.
    batch_size: (Optional) int. If the input is 2D, this might be the batch size
      of the 3D version of the `from_tensor` and `to_tensor`.
    always_attend_to_first_position: Should all blocks be able to attend to the
      `to_tensor`'s first position (e.g. a [CLS] position)?
    first_position_attends_to_all: Should the `from_tensor`'s first position
      be able to attend to all positions within the `from_tensor`?
    attend_from_chunk_width: The width of each block-wise chunk in
      `from_tensor`.
    attend_from_chunk_stride: The number of elements to skip when moving to the
      next block in `from_tensor`.
    attend_to_chunk_width: The width of each block-wise chunk in `to_tensor`.
    attend_to_chunk_stride: The number of elements to skip when moving to the
      next block in `to_tensor`.

  Returns:
    float Tensor of shape [batch_size, from_seq_length, num_attention_heads,
      size_per_head].

  Raises:
    ValueError: Any of the arguments or tensor shapes are invalid.
  
language/canine/local_attention.py