code: def pad_to_multiple_2d(x, block_shape):
    
    old_shape = x.get_shape().dims
    last = old_shape[-1]
    if len(old_shape) == 4:
      height_padding = -common_layers.shape_list(x)[1] % block_shape[0]
      width_padding = -common_layers.shape_list(x)[2] % block_shape[1]
      paddings = [[0, 0], [0, height_padding], [0, width_padding], [0, 0]]
    elif len(old_shape) == 5:
      height_padding = -common_layers.shape_list(x)[2] % block_shape[0]
      width_padding = -common_layers.shape_list(x)[3] % block_shape[1]
      paddings = [[0, 0], [0, 0], [0, height_padding], [0, width_padding], [0, 0]]
  
    padded_x = tf.pad(x, paddings)
    padded_shape = padded_x.get_shape().as_list()
    padded_shape = padded_shape[:-1] + [last]
    padded_x.set_shape(padded_shape)
    return padded_x
docs: Making sure x is a multiple of shape.
  
    Args:
      x: a [batch, heads, h, w, depth] or [batch, h, w, depth] tensor
      block_shape: a 2-d list of integer shapes
  
    Returns:
      padded_x: a [batch, heads, h, w, depth] or [batch, h, w, depth] tensor