def _scan_step_fn(state, example, packed_length, queue_size, spacing,
                  num_sequences, token_dtype):  # pylint: disable=g-doc-args
  """Transform function used by tf.data.experimental.scan to process an example.

  This is written as a stateless function rather than a class method because we
  trace it with AutoGraph (in order to simplify the conditional), and this way
  we don't have to worry about handling re-tracing semantics.

  Args:
    See the SequenceDatasetPacker class.

  Returns:
    The updated queue state, and either a packed example or a dummy sequence
    which will be filtered out downstream.
  """

  # Convert TensorArray tuples to lists since we'll need to replace them.
  availability, contents, top_index = state

  lengths = tf.concat([tf.shape(i) for i in example], axis=0)
  start_availability = availability.stack()
  can_fit = tf.reduce_all(tf.greater_equal(start_availability, lengths), axis=1)
  any_can_fit = tf.reduce_any(can_fit, axis=0)

  # AutoGraph will convert this block to a tf.cond
  if any_can_fit:
    # This indicates where in the FFD queue rotation a given index sits
    shifted_range = (
        tf.range(queue_size, dtype=INDEX_DTYPE) - top_index) % queue_size

    # Mark any indices which cannot accommodate the current example.
    exclusion_mask = tf.cast(tf.logical_not(can_fit), INDEX_DTYPE) * queue_size

    # Index in [0, queue_size) in which to place the sample. Note, this index
    # is the position in the actual TensorArray, not the index of the FFD queue.
    queue_index = (tf.reduce_min(shifted_range + exclusion_mask) +
                   top_index) % queue_size

    # NOTE(taylorrobie): We emit a non-empty Tensor for downstream checks.
    output_contents = -tf.ones((1, num_sequences), dtype=token_dtype)

  else:
    index_range = top_index * packed_length + tf.range(packed_length)
    output_contents = contents.gather(index_range)

    # Reset the queue state.
    availability = availability.write(
        top_index, packed_length * tf.ones((num_sequences,), dtype=INDEX_DTYPE))
    empty_contents = tf.zeros((packed_length, num_sequences * 2),
                              dtype=token_dtype)
    contents = contents.scatter(index_range, empty_contents)

    queue_index = top_index
    top_index = (top_index + 1) % queue_size

  pre_assign_availability = availability.read(queue_index)
  space_left = pre_assign_availability - lengths - spacing
  availability = availability.write(queue_index, space_left)

  # ============================================================================
  # == Update contents =========================================================
  # ============================================================================
  # Consider the following case for a seq-to-seq packing:
  #   (padding is represented as underscores)
  #
  #   Queue starting state:
  #     [1, 3, 2, 4, 6, 1, _, _, _, _, _, ...]
  #     [5, 9, _, _, _, _, _, _, _, _, _, ...]
  #
  #   Examples:
  #     [4, 2, 4], [3]
  #
  #   Desired new queue state:
  #     [1, 3, 2, 4, 6, 1, _, _, 4, 2, 4, _, _, ...]
  #     [5, 9, _, _, 3, _, _, _, _, _, _, _, _, ...]
  #
  # This could be acomplished by creating a TensorArray for each of the two
  # sequences, and scattering into the respective arrays. However TensorArray
  # writes are extremely expensive relative to other operations. So instead we
  # store the contents in a single TensorArray of shape (packed_length, 2), and
  # we pad and concatenate the examples such that they can be added in a single
  # assign:
  #
  #              [_, _, _, _, 4, 2, 4]
  #              [3, _, _, _, _, _, _]
  #                        +
  #  [1, 3, 2, 4, 6, 1, _, _, _, _, _, ...]
  #  [5, 9, _, _, _, _, _, _, _, _, _, ...]
  #
  # And in practice, the extra work of padding is neglidgable compared to
  # the gain from vectorizing the TensorArray assign. We also store a bit mask
  # denoting where sequences start which is used to compute segment and
  # position metadata:
  #
  #              [_, _, _, _, 1, _, _]
  #              [1, _, _, _, _, _, _]
  #                        +
  #  [1, _, _, _, _, _, _, _, _, _, _, ...]
  #  [1, _, _, _, _, _, _, _, _, _, _, ...]
  #
  # Both the contents and the mask are concatenated in the same TensorArray
  # for performance.

  start_index = packed_length - pre_assign_availability
  end_index = start_index + lengths
  leftmost = tf.reduce_min(start_index, axis=0)
  rightmost = tf.reduce_max(end_index, axis=0)
  delta = rightmost - leftmost
  pad_indices = [tf.stack((start_index[i] - leftmost, rightmost - end_index[i]))
                 for i in range(num_sequences)]

  padded_examples = [tf.pad(ex, padding[tf.newaxis, :])
                     for ex, padding in zip(example, pad_indices)]
  padded_examples = tf.transpose(tf.stack(padded_examples))
  mask_update = tf.one_hot(start_index - leftmost, delta,
                           dtype=contents.dtype, axis=0)

  content_update = tf.concat([padded_examples, mask_update], axis=1)

  index_range = (queue_index * packed_length +  # Offset into the right section.
                 tf.range(delta, dtype=INDEX_DTYPE) + leftmost)
  contents = contents.scatter(index_range, contents.gather(index_range) +
                              content_update)

  state = (availability, contents, top_index)
  return state, (tf.logical_not(any_can_fit), output_contents)
Transform function used by tf.data.experimental.scan to process an example.

  This is written as a stateless function rather than a class method because we
  trace it with AutoGraph (in order to simplify the conditional), and this way
  we don't have to worry about handling re-tracing semantics.

  Args:
    See the SequenceDatasetPacker class.

  Returns:
    The updated queue state, and either a packed example or a dummy sequence
    which will be filtered out downstream.
  
tensor2tensor/data_generators/generator_utils.py