def chunk_sequence(sequence, chunk_length=200, padding_value=0):
  """Split a nested dict of sequence tensors into a batch of chunks.

  This function does not expect a batch of sequences, but a single sequence. A
  `length` key is added if it did not exist already.

  Args:
    sequence: Nested dict of tensors with time dimension.
    chunk_length: Size of chunks the sequence will be split into.
    padding_value: Value used for padding the last chunk after the sequence.

  Returns:
    Nested dict of sequence tensors with chunk dimension.
  """
  if 'length' in sequence:
    length = sequence.pop('length')
  else:
    length = tf.shape(tools.nested.flatten(sequence)[0])[0]
  num_chunks = (length - 1) // chunk_length + 1
  padding_length = chunk_length * num_chunks - length
  padded = tools.nested.map(
      # pylint: disable=g-long-lambda
      lambda tensor: tf.concat([
          tensor, 0 * tensor[:padding_length] + padding_value], 0),
      sequence)
  chunks = tools.nested.map(
      # pylint: disable=g-long-lambda
      lambda tensor: tf.reshape(
          tensor, [num_chunks, chunk_length] + tensor.shape[1:].as_list()),
      padded)
  chunks['length'] = tf.concat([
      chunk_length * tf.ones((num_chunks - 1,), dtype=tf.int32),
      [chunk_length - padding_length]], 0)
  return chunks
Split a nested dict of sequence tensors into a batch of chunks.

  This function does not expect a batch of sequences, but a single sequence. A
  `length` key is added if it did not exist already.

  Args:
    sequence: Nested dict of tensors with time dimension.
    chunk_length: Size of chunks the sequence will be split into.
    padding_value: Value used for padding the last chunk after the sequence.

  Returns:
    Nested dict of sequence tensors with chunk dimension.