Dilated attention with a masking strategy.

def get_dilated_1d_attention_mask(
    num_heads, block_size,
    num_blocks, memory_size, gap_size,
    name="dilated_mask"):
  """Dilated attention with a masking strategy."""
  mask = np.ones((num_heads, block_size, 2*block_size), np.bool)

  # now going over every row to do the right assignment of
  # memory blocks
  for i in range(block_size):
    visible = 2*block_size  - (block_size-i)
    # You always attend to yourself, set the mask for that
    mask[:, i, -(block_size - i)] = 0
    # Maybe num_blocks can be automatically calculated?
    for j in range(num_blocks):
      for k in range(memory_size):
        index = ((gap_size + memory_size)*j) + k
        if index >= visible:
          break
        mask[:, i, -(index + block_size - i + 1)] = 0  # Verify

  # adding a num blocks dimension
  mask = np.expand_dims(mask, axis=1)
  return tf.constant(mask, dtype=tf.int32, name=name)