Pads unknown features' dimensions for TPU.

def pad_for_tpu(shapes_dict, hparams, max_length):
  """Pads unknown features' dimensions for TPU."""
  padded_shapes = {}

  def get_filler(specified_max_length):
    if not specified_max_length:
      return max_length
    return min(specified_max_length, max_length)

  inputs_none_filler = get_filler(hparams.max_input_seq_length)
  targets_none_filler = get_filler(hparams.max_target_seq_length)

  def pad_one_shape(shape, none_filler):
    return [
        (dim if dim is not None else none_filler) for dim in shape.as_list()
    ]

  for key, shape in six.iteritems(shapes_dict):
    if key == "inputs":
      padded_shapes[key] = pad_one_shape(shape, inputs_none_filler)
    elif key == "targets":
      padded_shapes[key] = pad_one_shape(shape, targets_none_filler)
    else:
      padded_shapes[key] = pad_one_shape(shape, max_length)
  return padded_shapes