parameters: (self, times)
return_statement: return looked_up_state
docstring: Retrieve cached states for a batch of times.
docstring_summary: Retrieve cached states for a batch of times.
docstring_tokens: Retrieve cached states for a batch of times .
function: def _get_cached_states(self, times):
      """Retrieve cached states for a batch of times."""
      read_chunk_numbers = self._get_chunk_number(times)
      looked_up_state = list(self._cached_states.lookup(
          math_ops.cast(read_chunk_numbers, dtypes.int64)))
      looked_up_state = tuple(looked_up_state)
      # We need to special-case the first chunk in a series to explicitly rely on
      # the model's starting state so that gradients flow back to it. Otherwise it
      # would affect only initialization, and would not be read from or updated
      # during training. Not doing this also isolates that part of the graph,
      # leading to errors on model reload if there are trainable variables
      # affecting a model's start state.
      if self._input_statistics is not None:
        start_time = self._input_statistics.start_time
      else:
        start_time = 0
      set_to_start_state = math_ops.equal(read_chunk_numbers,
                                          self._get_chunk_number(start_time))
      new_states = []
      for start_state_value, cache_variable in zip(
          nest.flatten(
              math_utils.replicate_state(self._start_state,
                                         array_ops.shape(times)[0])),
          nest.flatten(looked_up_state)):
  
        new_states.append(
            array_ops.where(set_to_start_state, start_state_value,
                            cache_variable))
      looked_up_state = nest.pack_sequence_as(looked_up_state, new_states)
      return looked_up_state
function_tokens: def _get_cached_states ( self , times ) : read_chunk_numbers = self . _get_chunk_number ( times ) looked_up_state = list ( self . _cached_states . lookup ( math_ops . cast ( read_chunk_numbers , dtypes . int64 ) ) ) looked_up_state = tuple ( looked_up_state ) # We need to special-case the first chunk in a series to explicitly rely on # the model's starting state so that gradients flow back to it. Otherwise it # would affect only initialization, and would not be read from or updated # during training. Not doing this also isolates that part of the graph, # leading to errors on model reload if there are trainable variables # affecting a model's start state. if self . _input_statistics is not None : start_time = self . _input_statistics . start_time else : start_time = 0 set_to_start_state = math_ops . equal ( read_chunk_numbers , self . _get_chunk_number ( start_time ) ) new_states = [ ] for start_state_value , cache_variable in zip ( nest . flatten ( math_utils . replicate_state ( self . _start_state , array_ops . shape ( times ) [ 0 ] ) ) , nest . flatten ( looked_up_state ) ) : new_states . append ( array_ops . where ( set_to_start_state , start_state_value , cache_variable ) ) looked_up_state = nest . pack_sequence_as ( looked_up_state , new_states ) return looked_up_state