parameters: (self, times, state, state_update_fn, outputs)
return_statement: return full_state, outputs_dict
docstring: Iterates over `times`, calling `state_update_fn` to collect outputs.
  
      Args:
        times: A [batch size x window size] Tensor of integers to iterate over.
        state: A list of model-specific state Tensors, each with shape [batch size
            x ...].
        state_update_fn: A callback taking the following arguments
              step_number; A scalar integer Tensor indicating the current position
                in the window.
              current_times; A [batch size] vector of Integers indicating times
                for each part of the batch.
              state; Current model state.
            It returns a tuple of (updated state, output_values), output_values
            being a dictionary of Tensors with keys corresponding to `outputs`.
        outputs: A list of strings indicating values which will be saved while
            iterating. Must match the keys of the dictionary returned by
            state_update_fn.
      Returns:
        A tuple of (state, output_dict)
        state: The final model state.
        output_dict: A dictionary of outputs corresponding to those specified in
          `outputs` and computed in state_update_fn.
docstring_summary: Iterates over `times`, calling `state_update_fn` to collect outputs.
docstring_tokens: Iterates over times calling state_update_fn to collect outputs .
function: def _state_update_loop(self, times, state, state_update_fn, outputs):
      """Iterates over `times`, calling `state_update_fn` to collect outputs.
  
      Args:
        times: A [batch size x window size] Tensor of integers to iterate over.
        state: A list of model-specific state Tensors, each with shape [batch size
            x ...].
        state_update_fn: A callback taking the following arguments
              step_number; A scalar integer Tensor indicating the current position
                in the window.
              current_times; A [batch size] vector of Integers indicating times
                for each part of the batch.
              state; Current model state.
            It returns a tuple of (updated state, output_values), output_values
            being a dictionary of Tensors with keys corresponding to `outputs`.
        outputs: A list of strings indicating values which will be saved while
            iterating. Must match the keys of the dictionary returned by
            state_update_fn.
      Returns:
        A tuple of (state, output_dict)
        state: The final model state.
        output_dict: A dictionary of outputs corresponding to those specified in
          `outputs` and computed in state_update_fn.
      """
      times = ops.convert_to_tensor(times, dtype=dtypes.int64)
      window_static_shape = times.get_shape()[1].value
      if self._static_unrolling_window_size_threshold is None:
        static_unroll = False
      else:
        # The user has specified a threshold for static loop unrolling.
        if window_static_shape is None:
          # We don't have static shape information for the window size, so dynamic
          # looping is our only option.
          static_unroll = False
        elif window_static_shape <= self._static_unrolling_window_size_threshold:
          # The threshold is satisfied; unroll statically
          static_unroll = True
        else:
          # A threshold was set but not satisfied
          static_unroll = False
  
      self._window_initializer(times, state)
  
      def _run_condition(step_number, *unused):
        del unused  # not part of while loop run condition
        return math_ops.less(step_number, window_size)
  
      def _state_update_step(
          step_number, state, state_accumulators, output_accumulators,
          reuse=False):
        """Impute, then take one state_update_fn step, accumulating outputs."""
        with variable_scope.variable_scope("state_update_step", reuse=reuse):
          current_times = times[:, step_number]
          state = self._imputation_step(current_times=current_times, state=state)
          output_accumulators_dict = {
              accumulator_key: accumulator
              for accumulator_key, accumulator
              in zip(outputs, output_accumulators)}
          step_state, output_values = state_update_fn(
              step_number=step_number,
              current_times=current_times,
              state=state)
          assert set(output_values.keys()) == set(outputs)
          new_output_accumulators = []
          for output_key in outputs:
            accumulator = output_accumulators_dict[output_key]
            output_value = output_values[output_key]
            new_output_accumulators.append(
                accumulator.write(step_number, output_value))
          flat_step_state = nest.flatten(step_state)
          assert len(state_accumulators) == len(flat_step_state)
          new_state_accumulators = []
          new_state_flat = []
          for step_state_value, state_accumulator, original_state in zip(
              flat_step_state, state_accumulators, nest.flatten(state)):
            # Make sure the static shape information is complete so while_loop
            # does not complain about shape information changing.
            step_state_value.set_shape(original_state.get_shape())
            new_state_flat.append(step_state_value)
            new_state_accumulators.append(state_accumulator.write(
                step_number, step_state_value))
          step_state = nest.pack_sequence_as(state, new_state_flat)
          return (step_number + 1, step_state,
                  new_state_accumulators, new_output_accumulators)
  
      window_size = array_ops.shape(times)[1]
  
      def _window_size_tensor_array(dtype):
        if static_unroll:
          return self._FakeTensorArray()
        else:
          return tensor_array_ops.TensorArray(
              dtype=dtype, size=window_size, dynamic_size=False)
  
      initial_loop_arguments = [
          array_ops.zeros([], dtypes.int32),
          state,
          [_window_size_tensor_array(element.dtype)
           for element in nest.flatten(state)],
          [_window_size_tensor_array(self.dtype) for _ in outputs]]
      if static_unroll:
        arguments = initial_loop_arguments
        for step_number in range(times.get_shape()[1].value):
          arguments = _state_update_step(
              array_ops.constant(step_number, dtypes.int32), *arguments[1:],
              reuse=(step_number > 0))  # Variable sharing between steps
      else:
        arguments = control_flow_ops.while_loop(
            cond=_run_condition,
            body=_state_update_step,
            loop_vars=initial_loop_arguments)
      (_, _, state_loop_result, outputs_loop_result) = arguments
  
      def _stack_and_transpose(tensor_array):
        """Stack and re-order the dimensions of a TensorArray."""
        if static_unroll:
          return array_ops.stack(tensor_array.values, axis=1)
        else:
          # TensorArrays from while_loop stack with window size as the first
          # dimension, so this function swaps it and the batch dimension to
          # maintain the [batch x window size x ...] convention used elsewhere.
          stacked = tensor_array.stack()
          return array_ops.transpose(
              stacked,
              perm=array_ops.concat([[1, 0], math_ops.range(
                  2, array_ops.rank(stacked))], 0))
  
      outputs_dict = {output_key: _stack_and_transpose(output)
                      for output_key, output
                      in zip(outputs, outputs_loop_result)}
      full_state = nest.pack_sequence_as(
          state,
          [_stack_and_transpose(state_element)
           for state_element in state_loop_result])
      return full_state, outputs_dict
function_tokens: def _state_update_loop ( self , times , state , state_update_fn , outputs ) : times = ops . convert_to_tensor ( times , dtype = dtypes . int64 ) window_static_shape = times . get_shape ( ) [ 1 ] . value if self . _static_unrolling_window_size_threshold is None : static_unroll = False else : # The user has specified a threshold for static loop unrolling. if window_static_shape is None : # We don't have static shape information for the window size, so dynamic # looping is our only option. static_unroll = False elif window_static_shape <= self . _static_unrolling_window_size_threshold : # The threshold is satisfied; unroll statically static_unroll = True else : # A threshold was set but not satisfied static_unroll = False self . _window_initializer ( times , state ) def _run_condition ( step_number , * unused ) : del unused # not part of while loop run condition return math_ops . less ( step_number , window_size ) def _state_update_step ( step_number , state , state_accumulators , output_accumulators , reuse = False ) : """Impute, then take one state_update_fn step, accumulating outputs.""" with variable_scope . variable_scope ( "state_update_step" , reuse = reuse ) : current_times = times [ : , step_number ] state = self . _imputation_step ( current_times = current_times , state = state ) output_accumulators_dict = { accumulator_key : accumulator for accumulator_key , accumulator in zip ( outputs , output_accumulators ) } step_state , output_values = state_update_fn ( step_number = step_number , current_times = current_times , state = state ) assert set ( output_values . keys ( ) ) == set ( outputs ) new_output_accumulators = [ ] for output_key in outputs : accumulator = output_accumulators_dict [ output_key ] output_value = output_values [ output_key ] new_output_accumulators . append ( accumulator . write ( step_number , output_value ) ) flat_step_state = nest . flatten ( step_state ) assert len ( state_accumulators ) == len ( flat_step_state ) new_state_accumulators = [ ] new_state_flat = [ ] for step_state_value , state_accumulator , original_state in zip ( flat_step_state , state_accumulators , nest . flatten ( state ) ) : # Make sure the static shape information is complete so while_loop # does not complain about shape information changing. step_state_value . set_shape ( original_state . get_shape ( ) ) new_state_flat . append ( step_state_value ) new_state_accumulators . append ( state_accumulator . write ( step_number , step_state_value ) ) step_state = nest . pack_sequence_as ( state , new_state_flat ) return ( step_number + 1 , step_state , new_state_accumulators , new_output_accumulators ) window_size = array_ops . shape ( times ) [ 1 ] def _window_size_tensor_array ( dtype ) : if static_unroll : return self . _FakeTensorArray ( ) else : return tensor_array_ops . TensorArray ( dtype = dtype , size = window_size , dynamic_size = False ) initial_loop_arguments = [ array_ops . zeros ( [ ] , dtypes . int32 ) , state , [ _window_size_tensor_array ( element . dtype ) for element in nest . flatten ( state ) ] , [ _window_size_tensor_array ( self . dtype ) for _ in outputs ] ] if static_unroll : arguments = initial_loop_arguments for step_number in range ( times . get_shape ( ) [ 1 ] . value ) : arguments = _state_update_step ( array_ops . constant ( step_number , dtypes . int32 ) , * arguments [ 1 : ] , reuse = ( step_number > 0 ) ) # Variable sharing between steps else : arguments = control_flow_ops . while_loop ( cond = _run_condition , body = _state_update_step , loop_vars = initial_loop_arguments ) ( _ , _ , state_loop_result , outputs_loop_result ) = arguments def _stack_and_transpose ( tensor_array ) : """Stack and re-order the dimensions of a TensorArray.""" if static_unroll : return array_ops . stack ( tensor_array . values , axis = 1 ) else : # TensorArrays from while_loop stack with window size as the first # dimension, so this function swaps it and the batch dimension to # maintain the [batch x window size x ...] convention used elsewhere. stacked = tensor_array . stack ( ) return array_ops . transpose ( stacked , perm = array_ops . concat ( [ [ 1 , 0 ] , math_ops . range ( 2 , array_ops . rank ( stacked ) ) ] , 0 ) ) outputs_dict = { output_key : _stack_and_transpose ( output ) for output_key , output in zip ( outputs , outputs_loop_result ) } full_state = nest . pack_sequence_as ( state , [ _stack_and_transpose ( state_element ) for state_element in state_loop_result ] ) return full_state , outputs_dict