def get_padded_shape(shape):
  """Return a `TensorShapeProto` with dims padded.

  The padding is based on the alignment requirement stored in `shape`.
  """
  alignment = shape.alignment
  remainder = shape.dims[-1] % shape.alignment
  if alignment == 0 or remainder == 0:
    return shape
  shape.dims[-1] += alignment - remainder;
  return shape