parameters: (self, input_dataset, initial_state, scan_func)
docstring: See `scan()` for details.
docstring_summary: See `scan()` for details.
docstring_tokens: See scan () for details .
function: def __init__(self, input_dataset, initial_state, scan_func):
      """See `scan()` for details."""
      super(_ScanDataset, self).__init__()
      self._input_dataset = input_dataset
  
      with ops.name_scope("initial_state"):
        self._initial_state = nest.pack_sequence_as(initial_state, [
            ops.convert_to_tensor(t, name="component_%d" % i)
            for i, t in enumerate(nest.flatten(initial_state))
        ])
  
      # Compute initial values for the state shapes and types based on
      # the initial state. These will be refined by running
      # `tf_scan_func` one or more times below.
      self._state_shapes = nest.pack_sequence_as(
          self._initial_state,
          [t.shape for t in nest.flatten(self._initial_state)])
      self._state_types = nest.pack_sequence_as(
          self._initial_state,
          [t.dtype for t in nest.flatten(self._initial_state)])
  
      # Will be populated by calling `tf_scan_func`.
      self._output_shapes = None
      self._output_types = None
  
      # Iteratively rerun the scan function until reaching a fixed pont on
      # `self._state_shapes`.
      need_to_rerun = True
      while need_to_rerun:
  
        flat_state_shapes = nest.flatten(self._state_shapes)
        flat_state_types = nest.flatten(self._state_types)
  
        # Create a list in which `tf_scan_func` will store the s
        flat_new_state_shapes = []
  
        @function.Defun(
            *(flat_state_types + nest.flatten(input_dataset.output_types)))
        def tf_scan_func(*args):
          """A wrapper for Defun that facilitates shape inference."""
          # Pass in shape information from the state and input_dataset.
          for arg, shape in zip(
              args,
              flat_state_shapes + nest.flatten(input_dataset.output_shapes)):
            arg.set_shape(shape)
  
          pivot = len(flat_state_shapes)
          old_state = nest.pack_sequence_as(self._initial_state, args[:pivot])
          input_value = nest.pack_sequence_as(input_dataset.output_types,
                                              args[pivot:])
  
          ret = scan_func(old_state, input_value)
          if not isinstance(ret, collections.Sequence) or len(ret) != 2:
            raise TypeError("The scan function must return a pair comprising the "
                            "new state and the output value.")
          new_state, output_value = ret
  
          flat_new_state = [
              ops.convert_to_tensor(t) for t in nest.flatten(new_state)
          ]
          flat_output_value = [
              ops.convert_to_tensor(t) for t in nest.flatten(output_value)
          ]
  
          # Extract shape information from the returned values.
          flat_new_state_shapes.extend([t.shape for t in flat_new_state])
          self._output_shapes = nest.pack_sequence_as(
              output_value, [t.shape for t in flat_output_value])
  
          # Extract and validate type information from the returned values.
          for t, dtype in zip(flat_new_state, flat_state_types):
            if t.dtype != dtype:
              raise TypeError(
                  "The element types for the new state must match the initial "
                  "state. Expected %s; got %s." %
                  (self._state_types, nest.pack_sequence_as(
                      self._state_types, [t.dtype for t in flat_new_state])))
          self._output_types = nest.pack_sequence_as(
              output_value, [t.dtype for t in flat_output_value])
  
          return flat_new_state + flat_output_value
  
        # Use the private method that will execute `tf_scan_func` but delay
        # adding it to the graph in case we need to rerun the function.
        tf_scan_func._create_definition_if_needed()  # pylint: disable=protected-access
  
        weakened_state_shapes = [
            original.most_specific_compatible_shape(new)
            for original, new in zip(flat_state_shapes, flat_new_state_shapes)
        ]
  
        need_to_rerun = False
        for original_shape, weakened_shape in zip(flat_state_shapes,
                                                  weakened_state_shapes):
          if original_shape.ndims is not None and (
              weakened_shape.ndims is None or
              original_shape.as_list() != weakened_shape.as_list()):
            need_to_rerun = True
            break
  
        if need_to_rerun:
          # NOTE(mrry): `self._output_shapes` will be overwritten when we rerun
          # `tf_scan_func`.
          self._state_shapes = nest.pack_sequence_as(self._state_shapes,
                                                     weakened_state_shapes)
  
      self._scan_func = tf_scan_func
function_tokens: def __init__ ( self , input_dataset , initial_state , scan_func ) : super ( _ScanDataset , self ) . __init__ ( ) self . _input_dataset = input_dataset with ops . name_scope ( "initial_state" ) : self . _initial_state = nest . pack_sequence_as ( initial_state , [ ops . convert_to_tensor ( t , name = "component_%d" % i ) for i , t in enumerate ( nest . flatten ( initial_state ) ) ] ) # Compute initial values for the state shapes and types based on # the initial state. These will be refined by running # `tf_scan_func` one or more times below. self . _state_shapes = nest . pack_sequence_as ( self . _initial_state , [ t . shape for t in nest . flatten ( self . _initial_state ) ] ) self . _state_types = nest . pack_sequence_as ( self . _initial_state , [ t . dtype for t in nest . flatten ( self . _initial_state ) ] ) # Will be populated by calling `tf_scan_func`. self . _output_shapes = None self . _output_types = None # Iteratively rerun the scan function until reaching a fixed pont on # `self._state_shapes`. need_to_rerun = True while need_to_rerun : flat_state_shapes = nest . flatten ( self . _state_shapes ) flat_state_types = nest . flatten ( self . _state_types ) # Create a list in which `tf_scan_func` will store the s flat_new_state_shapes = [ ] @ function . Defun ( * ( flat_state_types + nest . flatten ( input_dataset . output_types ) ) ) def tf_scan_func ( * args ) : """A wrapper for Defun that facilitates shape inference.""" # Pass in shape information from the state and input_dataset. for arg , shape in zip ( args , flat_state_shapes + nest . flatten ( input_dataset . output_shapes ) ) : arg . set_shape ( shape ) pivot = len ( flat_state_shapes ) old_state = nest . pack_sequence_as ( self . _initial_state , args [ : pivot ] ) input_value = nest . pack_sequence_as ( input_dataset . output_types , args [ pivot : ] ) ret = scan_func ( old_state , input_value ) if not isinstance ( ret , collections . Sequence ) or len ( ret ) != 2 : raise TypeError ( "The scan function must return a pair comprising the " "new state and the output value." ) new_state , output_value = ret flat_new_state = [ ops . convert_to_tensor ( t ) for t in nest . flatten ( new_state ) ] flat_output_value = [ ops . convert_to_tensor ( t ) for t in nest . flatten ( output_value ) ] # Extract shape information from the returned values. flat_new_state_shapes . extend ( [ t . shape for t in flat_new_state ] ) self . _output_shapes = nest . pack_sequence_as ( output_value , [ t . shape for t in flat_output_value ] ) # Extract and validate type information from the returned values. for t , dtype in zip ( flat_new_state , flat_state_types ) : if t . dtype != dtype : raise TypeError ( "The element types for the new state must match the initial " "state. Expected %s; got %s." % ( self . _state_types , nest . pack_sequence_as ( self . _state_types , [ t . dtype for t in flat_new_state ] ) ) ) self . _output_types = nest . pack_sequence_as ( output_value , [ t . dtype for t in flat_output_value ] ) return flat_new_state + flat_output_value # Use the private method that will execute `tf_scan_func` but delay # adding it to the graph in case we need to rerun the function. tf_scan_func . _create_definition_if_needed ( ) # pylint: disable=protected-access weakened_state_shapes = [ original . most_specific_compatible_shape ( new ) for original , new in zip ( flat_state_shapes , flat_new_state_shapes ) ] need_to_rerun = False for original_shape , weakened_shape in zip ( flat_state_shapes , weakened_state_shapes ) : if original_shape . ndims is not None and ( weakened_shape . ndims is None or original_shape . as_list ( ) != weakened_shape . as_list ( ) ) : need_to_rerun = True break if need_to_rerun : # NOTE(mrry): `self._output_shapes` will be overwritten when we rerun # `tf_scan_func`. self . _state_shapes = nest . pack_sequence_as ( self . _state_shapes , weakened_state_shapes ) self . _scan_func = tf_scan_func