def downsample_attention_mask(self,
                                char_attention_mask: tf.Tensor,
                                downsampling_rate: int,
                                dim: int = -1) -> tf.Tensor:
    """Downsample one dimension of an attention mask."""
    perm = None
    if dim != -1:
      ndims = 3
      perm = list(range(ndims))
      # Swap desired dimension with last dimension at beginning/end of
      # function.
      perm[dim], perm[-1] = perm[-1], perm[dim]

    if perm is not None:
      char_attention_mask = tf.transpose(char_attention_mask, perm)

    # `poolable_char_mask`: <float>[batch, char_seq, char_seq, 1]
    poolable_char_mask = tf.expand_dims(char_attention_mask, axis=-1)

    # `poolable_char_mask`: <float>[batch, from_seq, to_seq, 1]
    pooled_molecule_mask = tf.nn.max_pool2d(
        input=poolable_char_mask,
        ksize=[1, downsampling_rate],
        strides=[1, downsampling_rate],
        padding="VALID")

    # `molecule_attention_mask`: <float>[batch, from_seq, to_seq]
    molecule_attention_mask = tf.squeeze(pooled_molecule_mask, axis=-1)

    if perm is not None:
      molecule_attention_mask = tf.transpose(molecule_attention_mask, perm)
    return molecule_attention_mask
Downsample one dimension of an attention mask.
language/canine/modeling.py