def pad_to_multiple_nd(x, block_shape):
  """Making sure x is a multiple of shape.

  Args:
    x: a [batch, d1, d2, ..., dn, depth] tensor
    block_shape: a n-d list of integers representing block shape

  Returns:
    padded x where each dimension is a multiple of corresponding block length.
  """
  shape = common_layers.shape_list(x)
  paddings = [-l % b for l, b in zip(shape[1:-1], block_shape)]
  return tf.pad(x, [[0, 0]] + [[0, p] for p in paddings] + [[0, 0]])
Making sure x is a multiple of shape.

  Args:
    x: a [batch, d1, d2, ..., dn, depth] tensor
    block_shape: a n-d list of integers representing block shape

  Returns:
    padded x where each dimension is a multiple of corresponding block length.
  
tensor2tensor/layers/common_attention.py