code: def right_shift_blockwise(x, query_shape, name=None):
    
    with tf.variable_scope(
        name, default_name="right_shift_blockwise", values=[x]):
      x_list_shape = x.get_shape().as_list()
      x_shape = common_layers.shape_list(x)
      # Add a dummy dimension for heads.
      x = tf.expand_dims(x, axis=1)
      x = pad_to_multiple_2d(x, query_shape)
      padded_x_shape = common_layers.shape_list(x)
      # Set up q blocks.
      x_indices = gather_indices_2d(x, query_shape, query_shape)
      x_new = get_shifted_center_blocks(x, x_indices)
  
      # Put representations back into original shapes.
      output = scatter_blocks_2d(x_new, x_indices, padded_x_shape)
      # Remove the dummy head dimension.
      output = tf.squeeze(output, axis=1)
      # Remove the padding if introduced.
      output = tf.slice(output, [0, 0, 0, 0], [-1, x_shape[1], x_shape[2], -1])
      output.set_shape(x_list_shape)
      return output
docs: Right shifts once in every block.
  
    Args:
      x: a tensor of shape [batch, height, width, depth]
      query_shape: A 2d tuple of ints
      name: a string
  
    Returns:
      output: a tensor of the same shape as x