def right_shift_blockwise_nd(x, block_shape):
  """Right shift once in every block.

  Args:
    x: a [batch, d1, d2, ..., dn, depth] tensor
    block_shape: a tuple (q1, q2, ..., qn) representing the block shape

  Returns:
    a [batch, d1, d2, ..., dn, depth] tensor, right shifted.
  """
  blocked_x = break_into_blocks_nd(x, block_shape)
  blocked_x_shape = common_layers.shape_list(blocked_x)
  blocked_x = tf.reshape(blocked_x,
                         [blocked_x_shape[0], -1, blocked_x_shape[-1]])
  padded_x = tf.pad(blocked_x, [[0, 0], [1, 0], [0, 0]])
  x = tf.slice(padded_x, [0, 0, 0],
               [-1, np.prod(blocked_x_shape[1:-1], dtype=np.int32), -1])
  x = tf.reshape(x, blocked_x_shape)
  return put_back_blocks_nd(x, block_shape)
Right shift once in every block.

  Args:
    x: a [batch, d1, d2, ..., dn, depth] tensor
    block_shape: a tuple (q1, q2, ..., qn) representing the block shape

  Returns:
    a [batch, d1, d2, ..., dn, depth] tensor, right shifted.
  
tensor2tensor/layers/common_attention.py