def FixateInputShape(input_params,
                     batch_size,
                     target_max_seqlen=None,
                     source_max_seqlen=None):
  """Adjusts the input_params so that it always produces fixed shaped output.

  Sets p.target_max_length to target_max_seqlen and p.source_max_length to
  source_max_seqlen if set. Only keep items in bucket_upper_bound that is <=
  max_seqlen where max_seqlen is source_max_seqlen if given; otherwise,
  target_max_seqlen.

  Args:
    input_params: The input params.
    batch_size: The input generator should always output the batch size.
    target_max_seqlen: Every batch should produce samples with this sequence
      length, and samples are padded to this length. Use
      input_params.bucket_upper_bound[-1] if not set.
    source_max_seqlen: Same effect as target_max_seqlen but for source if set.

  Returns:
    input_params itself.
  """
  p = input_params
  # Pad to fixed length, since otherwise the infeed queue can't be set up,
  # because the shape won't be known statically.
  #
  # Limit memory by throwing away large sequences.
  if not target_max_seqlen:
    assert p.bucket_upper_bound
    target_max_seqlen = p.bucket_upper_bound[-1]
  if source_max_seqlen:
    p.bucket_upper_bound = [
        x for x in p.bucket_upper_bound if x <= source_max_seqlen
    ]
    if not p.bucket_upper_bound:
      p.bucket_upper_bound = [source_max_seqlen]
  else:
    p.bucket_upper_bound = [
        x for x in p.bucket_upper_bound if x <= target_max_seqlen
    ]
    if not p.bucket_upper_bound:
      p.bucket_upper_bound = [target_max_seqlen]
  p.bucket_batch_limit = [batch_size] * len(p.bucket_upper_bound)
  p.pad_to_max_seq_length = True
  if source_max_seqlen:
    p.source_max_length = source_max_seqlen
  p.target_max_length = target_max_seqlen
  if hasattr(p, 'pad_and_set_target_shape'):
    p.pad_and_set_target_shape = True

  # Because every batch is padded to the max sequence length and has the same
  # batch size, they are shardable.
  p.remote.shardable_batch = True
  return p