def multihead_attention_nd(query_antecedent,
                           memory_antecedent,
                           total_key_depth,
                           total_value_depth,
                           output_depth,
                           num_heads,
                           query_shape,
                           memory_flange,
                           masked=False,
                           cache=None,
                           decode_step=None,
                           name=None):
  """n-d Multihead scaled-dot-product attention with in/output transformations.

  Args:
    query_antecedent: a Tensor with shape [batch, d1, ..., dn, depth_q] or
      [batch, 1, ..., 1, depth_q] if in fast decoding mode.
    memory_antecedent: a Tensor with shape [batch, d1, ..., dn, depth_m] or None
      for self attention.
    total_key_depth: an integer
    total_value_depth: an integer
    output_depth: an integer
    num_heads: an integer dividing total_key_depth and total_value_depth
    query_shape: an tuple indicating the dimensions of each query block.
    memory_flange: an integer indicating how much to look around a query block
      in each dimension
    masked: a boolean to specify whether to do masked or unmasked attention.
    cache: a dict like: {
      'k': [batch, num_heads, d1, ..., dn, depth_k // num_heads],
      'v': [batch, num_heads, d1, ..., dn, depth_v // num_heads]} Caller should
        initially pass zero tensors for `decode_step` == 0. This method will
        update cache and caller should pass the same cache in consecutive calls.
        This works for both GPU and TPU inference. Caller should pass the latest
        query via `query_antecedent`. `memory_antecedent` should be None in this
        case, since auto-regressive decoding only applies to self attention.
    decode_step: integer to pass in decoding mode. `cache` and `decode_step`
      should both be set in decoding mode. Caller can also pass an empty `cache`
      without `decode_step`, for this method to initialize the cache for future
      calls with `decode_step` > 0.
    name: an optional string

  Returns:
    A Tensor of shape [batch, d1, ..., dn, output_depth] or
    [batch, 1, ..., 1, output_depth] if decode_step is set.

  Raises:
    ValueError: if the key depth or value depth are not divisible by the
      number of attention heads.
  """
  if total_key_depth % num_heads != 0:
    raise ValueError("Key depth (%d) must be divisible by the number of "
                     "attention heads (%d)." % (total_key_depth, num_heads))
  if total_value_depth % num_heads != 0:
    raise ValueError("Value depth (%d) must be divisible by the number of "
                     "attention heads (%d)." % (total_value_depth, num_heads))
  # Validate decoding input params are sensible.
  if decode_step is not None:
    assert "k" in cache and "v" in cache
  if cache is not None:
    assert memory_antecedent is None

  with tf.variable_scope(
      name,
      default_name="multihead_attention_nd",
      values=[query_antecedent, memory_antecedent]):
    if decode_step is not None:
      latest_antecedent = query_antecedent
      q, latest_k, latest_v = compute_qkv(latest_antecedent, None,
                                          total_key_depth, total_value_depth)
      latest_k = split_heads_nd(latest_k, num_heads)
      latest_v = split_heads_nd(latest_v, num_heads)
      # put latest k and v into their correct position in cache.
      k = cache["k"]
      v = cache["v"]
      k = put_item_in_decode_step(k, latest_k, decode_step, query_shape)
      v = put_item_in_decode_step(v, latest_v, decode_step, query_shape)
      cache["k"] = k
      cache["v"] = v

    else:
      q, k, v = compute_qkv(query_antecedent, memory_antecedent,
                            total_key_depth, total_value_depth)
      k = split_heads_nd(k, num_heads)
      v = split_heads_nd(v, num_heads)
      if cache is not None:
        cache["k"] = k
        cache["v"] = v
    # after splitting, shape is [batch, heads, d1, ..., dn, depth]
    q = split_heads_nd(q, num_heads)
    key_depth_per_head = total_key_depth // num_heads
    q *= key_depth_per_head**-0.5
    if masked:
      x = masked_local_attention_nd(
          q,
          k,
          v,
          query_shape=query_shape,
          memory_flange=memory_flange,
          decode_step=decode_step)
    else:
      raise NotImplementedError(
          "Unmaked multihead attention nd is not implemented")
    x = combine_heads_nd(x)
    x = common_layers.dense(
        x, output_depth, use_bias=False, name="output_transform")
    return x
n-d Multihead scaled-dot-product attention with in/output transformations.

  Args:
    query_antecedent: a Tensor with shape [batch, d1, ..., dn, depth_q] or
      [batch, 1, ..., 1, depth_q] if in fast decoding mode.
    memory_antecedent: a Tensor with shape [batch, d1, ..., dn, depth_m] or None
      for self attention.
    total_key_depth: an integer
    total_value_depth: an integer
    output_depth: an integer
    num_heads: an integer dividing total_key_depth and total_value_depth
    query_shape: an tuple indicating the dimensions of each query block.
    memory_flange: an integer indicating how much to look around a query block
      in each dimension
    masked: a boolean to specify whether to do masked or unmasked attention.
    cache: a dict like: {
      'k': [batch, num_heads, d1, ..., dn, depth_k // num_heads],
      'v': [batch, num_heads, d1, ..., dn, depth_v // num_heads]} Caller should
        initially pass zero tensors for `decode_step` == 0. This method will
        update cache and caller should pass the same cache in consecutive calls.
        This works for both GPU and TPU inference. Caller should pass the latest
        query via `query_antecedent`. `memory_antecedent` should be None in this
        case, since auto-regressive decoding only applies to self attention.
    decode_step: integer to pass in decoding mode. `cache` and `decode_step`
      should both be set in decoding mode. Caller can also pass an empty `cache`
      without `decode_step`, for this method to initialize the cache for future
      calls with `decode_step` > 0.
    name: an optional string

  Returns:
    A Tensor of shape [batch, d1, ..., dn, output_depth] or
    [batch, 1, ..., 1, output_depth] if decode_step is set.

  Raises:
    ValueError: if the key depth or value depth are not divisible by the
      number of attention heads.
  
tensor2tensor/layers/common_attention.py