def masked_local_attention_nd(q,
                              k,
                              v,
                              query_shape,
                              memory_flange,
                              decode_step=None,
                              name=None):
  """Masked local attention nd.

  Each position in q can attend to positions in memory that are positioned less
  than or equal to query position according to raster scan ordering and are in
  the same memory block. A memory block is n-dimensional and each dimension 'i'
  is of size q[i] + 2 * m[i] except for the first dimension which is of size
  q[0] + m[0]. NOTE: This computation assumes memory_flange is divisible by
  query_shape in every dimension.

  Args:
    q: a [batch, heads, d1, d2, ..., dn, depth_k] tensor or a [batch, heads, 1,
      1, ..., 1, depth_k] tensor in decoding mode.
    k: a [batch, heads, d1, d2, ..., dn, depth_k] tensor
    v: a [batch, heads, d1, d2, ..., dn, depth_v] tensor
    query_shape: a tuple (q1, q2, ..., qn) indicating the shape of query blocks.
    memory_flange: a tuple (m1, m2, ..., mn) indicating the number of extra
      positions in the attention memory. memory_shape=[q1 + m1, d2 + 2 * m2,
      ..., dn + 2 * mn]
    decode_step: an integer in fast decoding mode.
    name: an optional string

  Returns:
    a [batch, head, d1, d2, ..., dn, depth_v] tensor or
      [batch, head, 1, 1, ..., 1, depth_v] if decode_step is not None.
  """
  assert all([m % b == 0 for m, b in zip(memory_flange, query_shape)])
  with tf.variable_scope(
      name, default_name="masked_local_attention_nd", values=[q, k, v]):
    # This computation only applies to self attention, so assert q, k and v have
    # the same dimensions.
    if decode_step is None:
      q.get_shape().assert_is_compatible_with(k.get_shape())
      q.get_shape()[:-1].assert_is_compatible_with(v.get_shape()[:-1])
    else:
      k.get_shape().assert_is_compatible_with(v.get_shape())

    # move heads to batch dimension. This is needed to reduce number of
    # dimensions as much as possible, since most ops support only up to 7
    # dimensions.
    q_shape = common_layers.shape_list(q)
    k_shape = common_layers.shape_list(k)
    v_shape = common_layers.shape_list(v)
    q = tf.reshape(q, [-1] + q_shape[2:])
    k = tf.reshape(k, [-1] + k_shape[2:])
    v = tf.reshape(v, [-1] + v_shape[2:])

    # Pad query, key, value to ensure multiple of corresponding lengths.
    if decode_step is None:
      # don't pad query in fast decoding mode. We only need to calculate self
      # attention for one position.
      q = pad_to_multiple_nd(q, query_shape)
    k = pad_to_multiple_nd(k, query_shape)
    v = pad_to_multiple_nd(v, query_shape)

    # extract query and memory blocks
    if decode_step is None:
      q = break_into_blocks_nd(q, query_shape)
    else:
      # in fast decoding, q has 1 block with 1 item in it
      # q shape will be [batch] + [1] * n + [1, depth] which is equivalent of
      # [batch, b1, b2, ..., bn, items_in_block, depth] where there is 1 block
      # and 1 item in that block
      q = tf.reshape(q, [-1] + [1] * (len(q_shape) - 3) + [q_shape[-1]])
    k = break_into_memory_blocks_nd(k, query_shape, memory_flange, masked=True)
    v = break_into_memory_blocks_nd(v, query_shape, memory_flange, masked=True)

    # extract just one block of k and v in fast decoding mode.
    if decode_step is not None:
      k = select_block_for_decode_step(k, decode_step, query_shape)
      v = select_block_for_decode_step(v, decode_step, query_shape)

    # flatten q, k and v to [batch, num_blocks, items_in_block, depth]
    q, blocks_per_dim = flatten_blocks_nd(q)
    k, _ = flatten_blocks_nd(k)
    v, _ = flatten_blocks_nd(v)

    # make attention bias for causal attention.
    causal_attn_bias = causal_attention_bias_nd(
        query_shape, memory_flange, decode_step=decode_step)
    padding_attn_bias = tf.expand_dims(
        embedding_to_padding(v[:1, :, :, :]) * -1e9, axis=-2)

    if decode_step is None:
      num_blocks = common_layers.shape_list(v)[1]
      causal_attn_bias = tf.tile(causal_attn_bias, [1, num_blocks, 1, 1])
      padding_attn_bias = tf.tile(
          padding_attn_bias,
          [1, 1, np.prod(query_shape, dtype=np.int32), 1])
    attn_bias = tf.minimum(causal_attn_bias, padding_attn_bias)

    # Calculate dot product attention
    output = dot_product_attention(
        q,
        k,
        v,
        attn_bias,
        dropout_rate=0.,
        name=name or "masked_local_nd",
        make_image_summary=False)

    # restructure the output from blocks ordering to the original ordering
    output = unflatten_blocks_nd(output, blocks_per_dim)
    if decode_step is None:
      # In fast decoding, output only contains one element, this is not needed.
      output = put_back_blocks_nd(output, query_shape)

    # bring back the heads dimension
    output_shape = common_layers.shape_list(output)
    output = tf.reshape(output, q_shape[:2] + output_shape[1:])
    if decode_step is None:
      # No padding is introduced in fast decoding, no need to do this.
      output_shape = common_layers.shape_list(output)
      output = tf.slice(output, [0] * len(output_shape),
                        [-1, -1] + q_shape[2:-1] + [-1])
    return output
Masked local attention nd.

  Each position in q can attend to positions in memory that are positioned less
  than or equal to query position according to raster scan ordering and are in
  the same memory block. A memory block is n-dimensional and each dimension 'i'
  is of size q[i] + 2 * m[i] except for the first dimension which is of size
  q[0] + m[0]. NOTE: This computation assumes memory_flange is divisible by
  query_shape in every dimension.

  Args:
    q: a [batch, heads, d1, d2, ..., dn, depth_k] tensor or a [batch, heads, 1,
      1, ..., 1, depth_k] tensor in decoding mode.
    k: a [batch, heads, d1, d2, ..., dn, depth_k] tensor
    v: a [batch, heads, d1, d2, ..., dn, depth_v] tensor
    query_shape: a tuple (q1, q2, ..., qn) indicating the shape of query blocks.
    memory_flange: a tuple (m1, m2, ..., mn) indicating the number of extra
      positions in the attention memory. memory_shape=[q1 + m1, d2 + 2 * m2,
      ..., dn + 2 * mn]
    decode_step: an integer in fast decoding mode.
    name: an optional string

  Returns:
    a [batch, head, d1, d2, ..., dn, depth_v] tensor or
      [batch, head, 1, 1, ..., 1, depth_v] if decode_step is not None.
  
tensor2tensor/layers/common_attention.py