def get_spmd_image_padding_size(params, original_image_shape):
  """Returns the padding size to make it divisible by num_parts.

  SPMD partitions along the height dimension, this function returns the padding
  size along this dim.

  Args:
    params: Dictionary for SSD parameters.
    original_image_shape: 3D Image shape (H, W, C). If conv_space_to_depth
      is applied, the image shape after this optimization.

  Returns:
    Integer denoting the padding to add second dimension.
  """
  padding_size = 0
  if 'num_partitions' in params and params['num_partitions'] > 1:
    # SPMD is not supported with transpose input.
    assert not params['transpose_input']
    assert len(original_image_shape) == 3
    height_dim_ind = 0
    part_dim_size = original_image_shape[height_dim_ind]
    left_over = part_dim_size % params['num_partitions']
    if left_over:
      padding_size = params['num_partitions'] - left_over
  return padding_size