code: # Copyright 2017 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Defines ways of splicing and re-arranging time series.
  
  This file provides methods for reading, parsing, and re-arranging a time
  series. The main departure from standard TensorFlow input pipelines is a focus
  on "chunking" a time series, i.e. slicing it into small contiguous windows which
  are then batched together for training, a form of truncated
  backpropagation. This typically provides a significant speedup compared to
  looping over the whole series sequentially, by exploiting data parallelism and
  by reducing redundant contributions to gradients (due to redundant information
  in the series itself).
  
  A series, consisting of times (an increasing vector of integers) and values (one
  or more floating point values for each time) along with any exogenous features,
  is stored either in memory or on disk in various formats (e.g. "one record per
  timestep" on disk, or as a dictionary of Numpy arrays in memory). The location
  and format is specified by configuring a `TimeSeriesReader` object
  (e.g. `NumpyReader`, `CSVReader`), which reads the data into the TensorFlow
  graph. A `TimeSeriesInputFn` object (typically `RandomWindowInputFn`) then
  performs windowing and batching.
  
  Time series are passed through this pipeline as dictionaries mapping feature
  names to their values. For training and evaluation, these require at minimum
  `TrainEvalFeatures.TIMES` (scalar integers, one per timestep) and
  `TrainEvalFeatures.VALUES` (may be either univariate or multivariate). Exogenous
  features may have any shape, but are likewise associated with a timestep. Times
  themselves need not be contiguous or regular (although smaller/fewer gaps are
  generally better), but each timestep must have all `VALUES` and any exogenous
  features (i.e. times may be missing, but given that a time is specified, every
  other feature must also be specified for that step; some models may support
  making exogenous updates conditional).
  
  The expected use case of a `TimeSeriesInputFn` is that it is first configured
  (for example setting a batch or window size) and passed a reader (a
  `TimeSeriesReader` object). The `TimeSeriesInputFn` can then be passed as the
  input_fn of an Estimator.
  
  For example, `RandomWindowInputFn` is useful for creating batches of random
  chunks of a series for training:
  
  ```
    # Read data in the default "time,value" CSV format with no header
    reader = input_pipeline.CSVReader(csv_file_name)
    # Set up windowing and batching for training
    train_input_fn = input_pipeline.RandomWindowInputFn(
        reader, batch_size=16, window_size=16)
    # Fit model parameters to data
    estimator.train(input_fn=train_input_fn, steps=150)
  ```
  
  `RandomWindowInputFn` is the primary tool for training and quantitative
  evaluation of time series. `WholeDatasetInputFn`, which reads a whole series
  into memory, is useful for qualitative evaluation and preparing to make
  predictions with `predict_continuation_input_fn`.
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import abc
  
  import numpy
  
  from tensorflow.contrib.timeseries.python.timeseries import feature_keys
  from tensorflow.contrib.timeseries.python.timeseries import model_utils
  
  from tensorflow.python.estimator import estimator_lib
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import io_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn
  from tensorflow.python.ops import parsing_ops
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops import state_ops
  from tensorflow.python.ops import tensor_array_ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.training import input as input_lib
  from tensorflow.python.training import training
  from tensorflow.python.util import nest
  
  
  def predict_continuation_input_fn(
      evaluation, steps=None, times=None, exogenous_features=None):
    """An Estimator input_fn for running predict() after evaluate().
  
    If the call to evaluate() we are making predictions based on had a batch_size
    greater than one, predictions will start after each of these windows
    (i.e. will have the same batch dimension).
  
    Args:
      evaluation: The dictionary returned by `Estimator.evaluate`, with keys
        FilteringResults.STATE_TUPLE and FilteringResults.TIMES.
      steps: The number of steps to predict (scalar), starting after the
        evaluation. If `times` is specified, `steps` must not be; one is required.
      times: A [batch_size x window_size] array of integers (not a Tensor)
        indicating times to make predictions for. These times must be after the
        corresponding evaluation. If `steps` is specified, `times` must not be;
        one is required. If the batch dimension is omitted, it is assumed to be 1.
      exogenous_features: Optional dictionary. If specified, indicates exogenous
        features for the model to use while making the predictions. Values must
        have shape [batch_size x window_size x ...], where `batch_size` matches
        the batch dimension used when creating `evaluation`, and `window_size` is
        either the `steps` argument or the `window_size` of the `times` argument
        (depending on which was specified).
    Returns:
      An `input_fn` suitable for passing to the `predict` function of a time
      series `Estimator`.
    Raises:
      ValueError: If `times` or `steps` are misspecified.
    """
    if exogenous_features is None:
      exogenous_features = {}
    predict_times = model_utils.canonicalize_times_or_steps_from_output(
        times=times, steps=steps, previous_model_output=evaluation)
    features = {
        feature_keys.PredictionFeatures.STATE_TUPLE:
            evaluation[feature_keys.FilteringResults.STATE_TUPLE],
        feature_keys.PredictionFeatures.TIMES:
            predict_times
    }
    features.update(exogenous_features)
    def _predict_input_fn():
      """An input_fn for predict()."""
      # Prevents infinite iteration with a constant output in an Estimator's
      # predict().
      limited_features = {}
      for key, values in features.items():
        limited_values = nest.map_structure(
            lambda value: training.limit_epochs(value, num_epochs=1), values)
        limited_features[key] = limited_values
      return (limited_features, None)
    return _predict_input_fn
  
  
  class TimeSeriesReader(object):
    """Reads from and parses a data source for a `TimeSeriesInputFn`.
  
    This class provides methods that read a few records (`read`) or the full data
    set at once (`read_full`), and returns them as dictionaries mapping feature
    names to feature Tensors. Please see note at the top of the file for the
    structure of these dictionaries. The output is generally chunked by a
    `TimeSeriesInputFn` before being passed to the model.
    """
  
    def check_dataset_size(self, minimum_dataset_size):
      """When possible, raises an error if the dataset is too small.
  
      This method allows TimeSeriesReaders to raise informative error messages if
      the user has selected a window size in their TimeSeriesInputFn which is
      larger than the dataset size. However, many TimeSeriesReaders will not have
      access to a dataset size, in which case they do not need to override this
      method.
  
      Args:
        minimum_dataset_size: The minimum number of records which should be
          contained in the dataset. Readers should attempt to raise an error when
          possible if an epoch of data contains fewer records.
      """
      pass
  
    @abc.abstractmethod
    def read(self):
      """Parses one or more records into a feature dictionary.
  
      This method is expected to be called by a `TimeSeriesInputFn` object, and is
      not for use with models directly.
  
      A `TimeSeriesReader` object reads multiple records at a single time for
      efficiency; the size of these batches is an implementation detail internal
      to the input pipeline. These records should generally be sequential,
      although some out-of-order records due to file wraparounds are expected and
      must be handled by callers.
  
      Returns:
        A dictionary mapping feature names to `Tensor` values, each with an
        arbitrary batch dimension (for efficiency) as their first dimension.
      """
      pass
  
    @abc.abstractmethod
    def read_full(self):
      """Return the full dataset.
  
      Largely for interactive use/plotting (or evaluation on small
      datasets). Generally not very efficient. Not recommended for training.
  
      Returns:
        Same return type as `read`, but with the full dataset rather than an
        arbitrary chunk of it. A dictionary mapping feature names to `Tensor`
        values, where the size of the first dimension of each `Tensor` is the
        number of samples in the entire dataset. These `Tensor`s should be
        constant across graph invocations, assuming that the underlying data
        remains constant. Current implementations re-read data on each graph
        invocation, although this may change in the future.
      """
      pass
  
  
  class NumpyReader(TimeSeriesReader):
    """A time series parser for feeding Numpy arrays to a `TimeSeriesInputFn`.
  
    Avoids embedding data in the graph as constants.
    """
  
    def __init__(self, data, read_num_records_hint=4096):
      """Numpy array input for a `TimeSeriesInputFn`.
  
      Args:
        data: A dictionary mapping feature names to Numpy arrays, with two
          possible shapes (requires keys `TrainEvalFeatures.TIMES` and
          `TrainEvalFeatures.VALUES`):
            Univariate; `TIMES` and `VALUES` are both vectors of shape [series
              length]
            Multivariate; `TIMES` is a vector of shape [series length], `VALUES`
              has shape [series length x number of features].
          In any case, `VALUES` and any exogenous features must have their shapes
          prefixed by the shape of the value corresponding to the `TIMES` key.
        read_num_records_hint: The maximum number of samples to read at one time,
          for efficiency.
      """
      self._features = _canonicalize_numpy_data(
          data, require_single_batch=True)
      self._read_num_records_hint = read_num_records_hint
  
    def check_dataset_size(self, minimum_dataset_size):
      """Raise an error if the dataset is too small."""
      dataset_size = self._features[feature_keys.TrainEvalFeatures.TIMES].shape[1]
      if dataset_size < minimum_dataset_size:
        raise ValueError(
            ("A TimeSeriesInputFn is configured to create windows of size {}, "
             "but only {} records were available in the dataset. Either decrease "
             "the window size or provide more records.").format(
                 minimum_dataset_size, dataset_size))
  
    def read(self):
      """Returns a large chunk of the Numpy arrays for later re-chunking."""
      # Remove the batch dimension from all features
      features = {key: numpy.squeeze(value, axis=0)
                  for key, value in self._features.items()}
      return estimator_lib.inputs.numpy_input_fn(
          x=features,
          # The first dimensions of features are the series length, since we have
          # removed the batch dimension above. We now pull out
          # self._read_num_records_hint steps of this single time series to pass
          # to the TimeSeriesInputFn.
          batch_size=self._read_num_records_hint,
          num_epochs=None,
          shuffle=False)()
  
    def read_full(self):
      """Returns `Tensor` versions of the full Numpy arrays."""
      features = estimator_lib.inputs.numpy_input_fn(
          x=self._features,
          batch_size=1,
          num_epochs=None,
          queue_capacity=2,  # Each queue element is a full copy of the dataset
          shuffle=False)()
      # TimeSeriesInputFn expect just a batch dimension
      return {feature_name: array_ops.squeeze(feature_value, axis=0)
              for feature_name, feature_value in features.items()}
  
  
  class ReaderBaseTimeSeriesParser(TimeSeriesReader):
    """Base for time series readers which wrap a `tf.ReaderBase`."""
  
    def __init__(self, filenames, read_num_records_hint=4096):
      """Configure the time series reader.
  
      Args:
        filenames: A string or list of strings indicating files to read records
          from.
        read_num_records_hint: When not reading a full dataset, indicates the
          number of records to transfer in a single chunk (for efficiency). The
          actual number transferred at one time may vary.
      """
      self._filenames = filenames
      self._read_num_records_hint = read_num_records_hint
  
    @abc.abstractmethod
    def _get_reader(self):
      """Get an instance of the tf.ReaderBase associated with this class."""
      pass
  
    @abc.abstractmethod
    def _process_records(self, lines):
      """Given string items, return a processed dictionary of Tensors.
  
      Args:
        lines: A 1-dimensional string Tensor, each representing a record to parse
          (source dependent, e.g. a line of a file, or a serialized protocol
          buffer).
  
      Returns:
        A dictionary mapping feature names to their values. The batch dimensions
        should match the length of `lines`.
      """
      pass
  
    def _get_filename_queue(self, epoch_limit):
      """Constructs a filename queue with an epoch limit.
  
      `epoch_limit` is intended as an error checking fallback to prevent a reader
      from infinitely looping in its requests for more work items if none are
      available in any file. It should be set high enough that it is never reached
      assuming at least one record exists in some file.
  
      Args:
        epoch_limit: The maximum number of times to read through the complete list
          of files before throwing an OutOfRangeError.
      Returns:
        A tuple of (filename_queue, epoch_limiter):
          filename_queue: A FIFOQueue with filename work items.
          epoch_limiter: The local variable used for epoch limitation. This should
            be set to zero before a reader is passed `filename_queue` in order to
            reset the epoch limiter's state.
      """
      epoch_limiter = variable_scope.variable(
          initial_value=constant_op.constant(0, dtype=dtypes.int64),
          name="epoch_limiter",
          trainable=False,
          collections=[ops.GraphKeys.LOCAL_VARIABLES])
      filenames_tensor = array_ops.reshape(
          ops.convert_to_tensor(self._filenames), [-1])
      # We can't rely on epoch_limiter being initialized, since queue runners are
      # started before local variables are initialized. Instead, we ignore epoch
      # limits before variable initialization. This means that prior to variable
      # initialization, a QueueRunner may cause a reader to enter an un-checked
      # infinite loop. However, as soon as local variables are initialized, we
      # will start incrementing and checking epoch_limiter, which will interrupt
      # any in-progress loops.
      conditional_count_up_to = control_flow_ops.cond(
          state_ops.is_variable_initialized(epoch_limiter),
          lambda: epoch_limiter.count_up_to(epoch_limit),
          lambda: constant_op.constant(0, dtype=dtypes.int64))
      with ops.control_dependencies([conditional_count_up_to]):
        filenames_tensor = array_ops.identity(filenames_tensor)
      filename_queue = input_lib.string_input_producer(
          filenames_tensor, shuffle=False, capacity=1)
      return filename_queue, epoch_limiter
  
    def read(self):
      """Reads a chunk of data from the `tf.ReaderBase` for later re-chunking."""
      # Assuming there is at least one item to be read among all of the files in
      # self._filenames, we will not need to go through more than
      # self._read_num_records_hint epochs to get a batch of
      # self._read_num_records_hint records. Setting this limit and resetting it
      # before each reader.read_up_to call prevents infinite looping when there
      # are no records available in any of the files.
      filename_queue, epoch_limiter = self._get_filename_queue(
          epoch_limit=self._read_num_records_hint)
      reader = self._get_reader()
      epoch_reset_op = state_ops.assign(epoch_limiter, 0)
      with ops.control_dependencies([epoch_reset_op]):
        _, records = reader.read_up_to(
            filename_queue, self._read_num_records_hint)
      return self._process_records(records)
  
    def read_full(self):
      """Reads a full epoch of data into memory."""
      reader = self._get_reader()
      # Set a hard limit of 2 epochs through self._filenames. If there are any
      # records available, we should only end up reading the first record in the
      # second epoch before exiting the while loop and subsequently resetting the
      # epoch limit. If there are no records available in any of the files, this
      # hard limit prevents the reader.read_up_to call from looping infinitely.
      filename_queue, epoch_limiter = self._get_filename_queue(epoch_limit=2)
      epoch_reset_op = state_ops.assign(epoch_limiter, 0)
      with ops.control_dependencies([epoch_reset_op]):
        first_key, first_value = reader.read_up_to(filename_queue, 1)
      # Read until we get a duplicate key (one epoch)
      def _while_condition(
          current_key, current_value, current_index, collected_records):
        del current_value, current_index, collected_records  # unused
        return math_ops.not_equal(array_ops.squeeze(current_key, axis=0),
                                  array_ops.squeeze(first_key, axis=0))
  
      def _while_body(
          current_key, current_value, current_index, collected_records):
        del current_key  # unused
        new_key, new_value = reader.read_up_to(filename_queue, 1)
        new_key.set_shape([1])
        new_value.set_shape([1])
        return (new_key,
                new_value,
                current_index + 1,
                collected_records.write(current_index, current_value))
      _, _, _, records_ta = control_flow_ops.while_loop(
          _while_condition,
          _while_body,
          [constant_op.constant([""]), first_value,
           0,  # current_index starting value
           tensor_array_ops.TensorArray(  # collected_records
               dtype=dtypes.string, size=0, dynamic_size=True)])
      records = records_ta.concat()
      # Reset the reader when we're done so that subsequent requests for data get
      # the dataset in the proper order.
      with ops.control_dependencies([records]):
        reader_reset_op = reader.reset()
      with ops.control_dependencies([reader_reset_op]):
        records = array_ops.identity(records)
      return self._process_records(records)
  
  
  class CSVReader(ReaderBaseTimeSeriesParser):
    """Reads from a collection of CSV-formatted files."""
  
    def __init__(self,
                 filenames,
                 column_names=(feature_keys.TrainEvalFeatures.TIMES,
                               feature_keys.TrainEvalFeatures.VALUES),
                 column_dtypes=None,
                 skip_header_lines=None,
                 read_num_records_hint=4096):
      """CSV-parsing reader for a `TimeSeriesInputFn`.
  
      Args:
        filenames: A filename or list of filenames to read the time series
            from. Each line must have columns corresponding to `column_names`.
        column_names: A list indicating names for each
            feature. `TrainEvalFeatures.TIMES` and `TrainEvalFeatures.VALUES` are
            required; `VALUES` may be repeated to indicate a multivariate series.
        column_dtypes: If provided, must be a list with the same length as
            `column_names`, indicating dtypes for each column. Defaults to
            `tf.int64` for `TrainEvalFeatures.TIMES` and `tf.float32` for
            everything else.
        skip_header_lines: Passed on to `tf.TextLineReader`; skips this number of
            lines at the beginning of each file.
        read_num_records_hint: When not reading a full dataset, indicates the
            number of records to parse/transfer in a single chunk (for
            efficiency). The actual number transferred at one time may be more or
            less.
      Raises:
        ValueError: If required column names are not specified, or if lengths do
          not match.
      """
      if feature_keys.TrainEvalFeatures.TIMES not in column_names:
        raise ValueError("'{}' is a required column.".format(
            feature_keys.TrainEvalFeatures.TIMES))
      if feature_keys.TrainEvalFeatures.VALUES not in column_names:
        raise ValueError("'{}' is a required column.".format(
            feature_keys.TrainEvalFeatures.VALUES))
      if column_dtypes is not None and len(column_dtypes) != len(column_names):
        raise ValueError(
            ("If specified, the length of column_dtypes must match the length of "
             "column_names (got column_dtypes={} and column_names={}).").format(
                 column_dtypes, column_names))
      if sum(1 for column_name in column_names
             if column_name == feature_keys.TrainEvalFeatures.TIMES) != 1:
        raise ValueError(
            "Got more than one times column ('{}'), but exactly "
            "one is required.".format(feature_keys.TrainEvalFeatures.TIMES))
      self._column_names = column_names
      self._column_dtypes = column_dtypes
      self._skip_header_lines = skip_header_lines
      super(CSVReader, self).__init__(
          filenames=filenames, read_num_records_hint=read_num_records_hint)
  
    def _get_reader(self):
      return io_ops.TextLineReader(skip_header_lines=self._skip_header_lines)
  
    def _process_records(self, lines):
      """Parse `lines` as CSV records."""
      if self._column_dtypes is None:
        default_values = [(array_ops.zeros([], dtypes.int64),)
                          if column_name == feature_keys.TrainEvalFeatures.TIMES
                          else () for column_name in self._column_names]
      else:
        default_values = [(array_ops.zeros([], dtype),)
                          for dtype in self._column_dtypes]
      columns = parsing_ops.decode_csv(lines, default_values)
      features_lists = {}
      for column_name, value in zip(self._column_names, columns):
        features_lists.setdefault(column_name, []).append(value)
      features = {}
      for column_name, values in features_lists.items():
        if column_name == feature_keys.TrainEvalFeatures.TIMES:
          features[column_name] = values[0]
        else:
          features[column_name] = array_ops.stack(values, axis=1)
      return features
  
  
  class TFExampleReader(ReaderBaseTimeSeriesParser):
    """Reads and parses `tf.Example`s from a TFRecords file."""
  
    def __init__(self,
                 filenames,
                 features):
      """Configure `tf.Example` parsing.
  
      Args:
        filenames: A filename or list of filenames to read the time series
            from. Each line must have columns corresponding to `column_names`.
        features: A dictionary mapping from feature keys to `tf.FixedLenFeature`
            objects. Must include `TrainEvalFeatures.TIMES` (scalar integer) and
            `TrainEvalFeatures.VALUES` (floating point vector) features.
      Raises:
        ValueError: If required times/values features are not present.
      """
      if feature_keys.TrainEvalFeatures.TIMES not in features:
        raise ValueError("'{}' is a required column.".format(
            feature_keys.TrainEvalFeatures.TIMES))
      if feature_keys.TrainEvalFeatures.VALUES not in features:
        raise ValueError("'{}' is a required column.".format(
            feature_keys.TrainEvalFeatures.VALUES))
      self._features = features
      super(TFExampleReader, self).__init__(filenames=filenames)
  
    def _get_reader(self):
      return io_ops.TFRecordReader()
  
    def _process_records(self, examples):
      """Parse `tf.Example`s into `Tensors`."""
      return parsing_ops.parse_example(
          serialized=examples, features=self._features)
  
  
  class TimeSeriesInputFn(object):
    """Base for classes which create batches of windows from a time series."""
  
    @abc.abstractmethod
    def create_batch(self):
      """Creates chunked Tensors from times, values, and other features.
  
      Suitable for use as the input_fn argument of a tf.estimator.Estimator's
      fit() or evaluate() method.
  
      Returns:
        A tuple of (features, targets):
          features: A dictionary with `TrainEvalFeatures.TIMES` and
            `TrainEvalFeatures.VALUES` as keys, `TIMES` having an associated value
            with shape [batch size x window length], `VALUES` with shape [batch
            size x window length x number of features]. Any other features will
            also have shapes prefixed with [batch size x window length].
          targets: Not used, but must have a value for compatibility with the
            Estimator API. That value should be None.
      """
      pass
  
    def __call__(self):
      # Allow a TimeSeriesInputFn to be used as an input function directly
      return self.create_batch()
  
  
  class WholeDatasetInputFn(TimeSeriesInputFn):
    """Supports passing a full time series to a model for evaluation/inference.
  
    Note that this `TimeSeriesInputFn` is not designed for high throughput, and
    should not be used for training. It allows for sequential evaluation on a full
    dataset (with sequential in-sample predictions), which then feeds naturally
    into `predict_continuation_input_fn` for making out-of-sample
    predictions. While this is useful for plotting and interactive use,
    `RandomWindowInputFn` is better suited to training and quantitative
    evaluation.
    """
    # TODO(allenl): A SequentialWindowInputFn for getting model end state without
    # loading the whole dataset into memory (or for quantitative evaluation of
    # sequential models). Note that an Estimator using such a TimeSeriesInputFn
    # won't return in-sample predictions for the whole dataset, which means it
    # won't be terribly useful for interactive use/plotting (unless the user
    # passes in concat metrics). Also need to be careful about state saving for
    # sequential models, particularly the gaps between chunks.
  
    def __init__(self, time_series_reader):
      """Initialize the `TimeSeriesInputFn`.
  
      Args:
        time_series_reader: A TimeSeriesReader object.
      """
      self._reader = time_series_reader
      super(WholeDatasetInputFn, self).__init__()
  
    def create_batch(self):
      """A suitable `input_fn` for an `Estimator`'s `evaluate()`.
  
      Returns:
        A dictionary mapping feature names to `Tensors`, each shape
        prefixed by [1, data set size] (i.e. a batch size of 1).
      """
      features = self._reader.read_full()
      # Add a batch dimension of one to each feature.
      return ({feature_name: feature_value[None, ...]
               for feature_name, feature_value in features.items()},
              None)
  
  
  class RandomWindowInputFn(TimeSeriesInputFn):
    """Wraps a `TimeSeriesReader` to create random batches of windows.
  
    Tensors are first collected into sequential windows (in a windowing queue
    created by `tf.train.batch`, based on the order returned from
    `time_series_reader`), then these windows are randomly batched (in a
    `RandomShuffleQueue`), the Tensors returned by `create_batch` having shapes
    prefixed by [`batch_size`, `window_size`].
  
    This `TimeSeriesInputFn` is useful for both training and quantitative
    evaluation (but be sure to run several epochs for sequential models such as
    `StructuralEnsembleRegressor` to completely flush stale state left over from
    training). For qualitative evaluation or when preparing for predictions, use
    `WholeDatasetInputFn`.
    """
  
    def __init__(
        self, time_series_reader, window_size, batch_size,
        queue_capacity_multiplier=1000, shuffle_min_after_dequeue_multiplier=2,
        discard_out_of_order=True, discard_consecutive_batches_limit=1000,
        jitter=True, num_threads=2, shuffle_seed=None):
      """Configure the RandomWindowInputFn.
  
      Args:
        time_series_reader: A TimeSeriesReader object.
        window_size: The number of examples to keep together sequentially. This
          controls the length of truncated backpropagation: smaller values mean
          less sequential computation, which can lead to faster training, but
          create a coarser approximation to the gradient (which would ideally be
          computed by a forward pass over the entire sequence in order).
        batch_size: The number of windows to place together in a batch. Larger
          values will lead to more stable gradients during training.
        queue_capacity_multiplier: The capacity for the queues used to create
          batches, specified as a multiple of `batch_size` (for
          RandomShuffleQueue) and `batch_size * window_size` (for the
          FIFOQueue). Controls the maximum number of windows stored. Should be
          greater than `shuffle_min_after_dequeue_multiplier`.
        shuffle_min_after_dequeue_multiplier: The minimum number of windows in the
          RandomShuffleQueue after a dequeue, which controls the amount of entropy
          introduced during batching. Specified as a multiple of `batch_size`.
        discard_out_of_order: If True, windows of data which have times which
          decrease (a higher time followed by a lower time) are discarded. If
          False, the window and associated features are instead sorted so that
          times are non-decreasing. Discarding is typically faster, as models do
          not have to deal with artificial gaps in the data. However, discarding
          does create a bias where the beginnings and endings of files are
          under-sampled.
        discard_consecutive_batches_limit: Raise an OutOfRangeError if more than
          this number of batches are discarded without a single non-discarded
          window (prevents infinite looping when the dataset is too small).
        jitter: If True, randomly discards examples between some windows in order
          to avoid deterministic chunking patterns. This is important for models
          like AR which may otherwise overfit a fixed chunking.
        num_threads: Use this number of threads for queues. Setting a value of 1
          removes one source of non-determinism (and in combination with
          shuffle_seed should provide deterministic windowing).
        shuffle_seed: A seed for window shuffling. The default value of None
          provides random behavior. With `shuffle_seed` set and
          `num_threads=1`, provides deterministic behavior.
      """
      self._reader = time_series_reader
      self._window_size = window_size
      self._reader.check_dataset_size(minimum_dataset_size=self._window_size)
      self._batch_size = batch_size
      self._queue_capacity_multiplier = queue_capacity_multiplier
      self._shuffle_min_after_dequeue_multiplier = (
          shuffle_min_after_dequeue_multiplier)
      self._discard_out_of_order = discard_out_of_order
      self._discard_limit = discard_consecutive_batches_limit
      self._jitter = jitter
      if num_threads is None:
        self._num_threads = self._batch_size
      else:
        self._num_threads = num_threads
      self._shuffle_seed = shuffle_seed
      super(RandomWindowInputFn, self).__init__()
  
    def create_batch(self):
      """Create queues to window and batch time series data.
  
      Returns:
        A dictionary of Tensors corresponding to the output of `self._reader`
        (from the `time_series_reader` constructor argument), each with shapes
        prefixed by [`batch_size`, `window_size`].
      """
      features = self._reader.read()
      if self._jitter:
        # TODO(agarwal, allenl): Figure out if more jitter is needed here.
        jitter = random_ops.random_uniform(shape=[], maxval=2, dtype=dtypes.int32)
      else:
        jitter = 0
      # To keep things efficient, we pass from the windowing batcher to the
      # batch-of-windows batcher in batches. This avoids the need for huge numbers
      # of threads, but does mean that jitter is only applied occasionally.
      # TODO(allenl): Experiment with different internal passing sizes.
      internal_passing_size = self._batch_size
      features_windowed = input_lib.batch(
          features,
          batch_size=self._window_size * internal_passing_size + jitter,
          enqueue_many=True,
          capacity=(self._queue_capacity_multiplier
                    * internal_passing_size * self._window_size),
          num_threads=self._num_threads)
      raw_features_windowed = features_windowed
      if self._jitter:
        features_windowed = {
            key: value[jitter:]
            for key, value in features_windowed.items()}
      features_windowed = {
          key: array_ops.reshape(
              value,
              array_ops.concat(
                  [[internal_passing_size, self._window_size],
                   array_ops.shape(value)[1:]],
                  axis=0))
          for key, value in features_windowed.items()}
      batch_and_window_shape = tensor_shape.TensorShape(
          [internal_passing_size, self._window_size])
      for key in features_windowed.keys():
        features_windowed[key].set_shape(
            batch_and_window_shape.concatenate(
                raw_features_windowed[key].get_shape()[1:]))
      # When switching files, we may end up with windows where the time is not
      # decreasing, even if times within each file are sorted (and even if those
      # files are visited in order, when looping back around to the beginning of
      # the first file). This is hard for models to deal with, so we either
      # discard such examples, creating a bias where the beginning and end of the
      # series is under-sampled, or we sort the window, creating large gaps.
      times = features_windowed[feature_keys.TrainEvalFeatures.TIMES]
      if self._discard_out_of_order:
        non_decreasing = math_ops.reduce_all(
            times[:, 1:] >= times[:, :-1], axis=1)
        # Ensure that no more than self._discard_limit complete batches are
        # discarded contiguously (resetting the count when we find a single clean
        # window). This prevents infinite looping when the dataset is smaller than
        # the window size.
        # TODO(allenl): Figure out a way to return informative errors from
        # count_up_to.
        discarded_windows_limiter = variable_scope.variable(
            initial_value=constant_op.constant(0, dtype=dtypes.int64),
            name="discarded_windows_limiter",
            trainable=False,
            collections=[ops.GraphKeys.LOCAL_VARIABLES])
        def _initialized_limit_check():
          return control_flow_ops.cond(
              math_ops.reduce_any(non_decreasing),
              lambda: state_ops.assign(discarded_windows_limiter, 0),
              lambda: discarded_windows_limiter.count_up_to(self._discard_limit))
        discard_limit_op = control_flow_ops.cond(
            state_ops.is_variable_initialized(discarded_windows_limiter),
            _initialized_limit_check,
            lambda: constant_op.constant(0, dtype=dtypes.int64))
        with ops.control_dependencies([discard_limit_op]):
          non_decreasing = array_ops.identity(non_decreasing)
      else:
        _, indices_descending = nn.top_k(
            times, k=array_ops.shape(times)[-1], sorted=True)
        indices = array_ops.reverse(indices_descending, axis=[0])
        features_windowed = {
            key: array_ops.gather(params=value, indices=indices)
            for key, value in features_windowed.items()
        }
        non_decreasing = True
      features_batched = input_lib.maybe_shuffle_batch(
          features_windowed,
          num_threads=self._num_threads,
          seed=self._shuffle_seed,
          batch_size=self._batch_size,
          capacity=self._queue_capacity_multiplier * self._batch_size,
          min_after_dequeue=(self._shuffle_min_after_dequeue_multiplier *
                             self._batch_size),
          keep_input=non_decreasing,
          enqueue_many=True)
      return (features_batched, None)
  
  
  def _canonicalize_numpy_data(data, require_single_batch):
    """Do basic checking and reshaping for Numpy data.
  
    Args:
      data: A dictionary mapping keys to Numpy arrays, with several possible
        shapes (requires keys `TrainEvalFeatures.TIMES` and
        `TrainEvalFeatures.VALUES`):
          Single example; `TIMES` is a scalar and `VALUES` is either a scalar or a
            vector of length [number of features].
          Sequence; `TIMES` is a vector of shape [series length], `VALUES` either
            has shape [series length] (univariate) or [series length x number of
            features] (multivariate).
          Batch of sequences; `TIMES` is a vector of shape [batch size x series
            length], `VALUES` has shape [batch size x series length] or [batch
            size x series length x number of features].
        In any case, `VALUES` and any exogenous features must have their shapes
        prefixed by the shape of the value corresponding to the `TIMES` key.
      require_single_batch: If True, raises an error if the provided data has a
        batch dimension > 1.
    Returns:
      A dictionary with features normalized to have shapes prefixed with [batch
      size x series length]. The sizes of dimensions which were omitted in the
      inputs are 1.
    Raises:
      ValueError: If dimensions are incorrect or do not match, or required
        features are missing.
    """
    features = {key: numpy.array(value) for key, value in data.items()}
    if (feature_keys.TrainEvalFeatures.TIMES not in features or
        feature_keys.TrainEvalFeatures.VALUES not in features):
      raise ValueError("{} and {} are required features.".format(
          feature_keys.TrainEvalFeatures.TIMES,
          feature_keys.TrainEvalFeatures.VALUES))
    times = features[feature_keys.TrainEvalFeatures.TIMES]
    for key, value in features.items():
      if value.shape[:len(times.shape)] != times.shape:
        raise ValueError(
            ("All features must have their shapes prefixed by the shape of the"
             " times feature. Got shape {} for feature '{}', but shape {} for"
             " '{}'").format(value.shape, key, times.shape,
                             feature_keys.TrainEvalFeatures.TIMES))
    if not times.shape:  # a single example
      if not features[feature_keys.TrainEvalFeatures.VALUES].shape:  # univariate
        # Add a feature dimension (with one feature)
        features[feature_keys.TrainEvalFeatures.VALUES] = features[
            feature_keys.TrainEvalFeatures.VALUES][..., None]
      elif len(features[feature_keys.TrainEvalFeatures.VALUES].shape) > 1:
        raise ValueError(
            ("Got an unexpected number of dimensions for the '{}' feature."
             " Was expecting at most 1 dimension"
             " ([number of features]) since '{}' does not "
             "have a batch or time dimension, but got shape {}").format(
                 feature_keys.TrainEvalFeatures.VALUES,
                 feature_keys.TrainEvalFeatures.TIMES,
                 features[feature_keys.TrainEvalFeatures.VALUES].shape))
      # Add trivial batch and time dimensions for every feature
      features = {key: value[None, None, ...] for key, value in features.items()}
    if len(times.shape) == 1:  # shape [series length]
      if len(features[feature_keys.TrainEvalFeatures.VALUES]
             .shape) == 1:  # shape [series length]
        # Add a feature dimension (with one feature)
        features[feature_keys.TrainEvalFeatures.VALUES] = features[
            feature_keys.TrainEvalFeatures.VALUES][..., None]
      elif len(features[feature_keys.TrainEvalFeatures.VALUES].shape) > 2:
        raise ValueError(
            ("Got an unexpected number of dimensions for the '{}' feature."
             " Was expecting at most 2 dimensions"
             " ([series length, number of features]) since '{}' does not "
             "have a batch dimension, but got shape {}").format(
                 feature_keys.TrainEvalFeatures.VALUES,
                 feature_keys.TrainEvalFeatures.TIMES,
                 features[feature_keys.TrainEvalFeatures.VALUES].shape))
      # Add trivial batch dimensions for every feature
      features = {key: value[None, ...] for key, value in features.items()}
    elif len(features[feature_keys.TrainEvalFeatures.TIMES]
             .shape) != 2:  # shape [batch size, series length]
      raise ValueError(
          ("Got an unexpected number of dimensions for times. Was expecting at "
           "most two ([batch size, series length]), but got shape {}.").format(
               times.shape))
    if require_single_batch:
      # We don't expect input to be already batched; batching is done later
      if features[feature_keys.TrainEvalFeatures.TIMES].shape[0] != 1:
        raise ValueError("Got batch input, was expecting unbatched input.")
    return features
   # Copyright 2015 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Operations for generating random numbers."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import random_seed
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import gen_random_ops
  from tensorflow.python.ops import math_ops
  # go/tf-wildcard-import
  # pylint: disable=wildcard-import
  from tensorflow.python.ops.gen_random_ops import *
  from tensorflow.python.util import deprecation
  from tensorflow.python.util.tf_export import tf_export
  
  # pylint: enable=wildcard-import
  
  
  def _ShapeTensor(shape):
    """Convert to an int32 or int64 tensor, defaulting to int32 if empty."""
    if isinstance(shape, (tuple, list)) and not shape:
      dtype = dtypes.int32
    else:
      dtype = None
    return ops.convert_to_tensor(shape, dtype=dtype, name="shape")
  
  
  @tf_export("random.normal", "random_normal")
  def random_normal(shape,
                    mean=0.0,
                    stddev=1.0,
                    dtype=dtypes.float32,
                    seed=None,
                    name=None):
    """Outputs random values from a normal distribution.
  
    Args:
      shape: A 1-D integer Tensor or Python array. The shape of the output tensor.
      mean: A 0-D Tensor or Python value of type `dtype`. The mean of the normal
        distribution.
      stddev: A 0-D Tensor or Python value of type `dtype`. The standard deviation
        of the normal distribution.
      dtype: The type of the output.
      seed: A Python integer. Used to create a random seed for the distribution.
        See
        `tf.set_random_seed`
        for behavior.
      name: A name for the operation (optional).
  
    Returns:
      A tensor of the specified shape filled with random normal values.
    """
    with ops.name_scope(name, "random_normal", [shape, mean, stddev]) as name:
      shape_tensor = _ShapeTensor(shape)
      mean_tensor = ops.convert_to_tensor(mean, dtype=dtype, name="mean")
      stddev_tensor = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev")
      seed1, seed2 = random_seed.get_seed(seed)
      rnd = gen_random_ops.random_standard_normal(
          shape_tensor, dtype, seed=seed1, seed2=seed2)
      mul = rnd * stddev_tensor
      value = math_ops.add(mul, mean_tensor, name=name)
      return value
  
  
  ops.NotDifferentiable("RandomStandardNormal")
  
  
  def parameterized_truncated_normal(shape,
                                     means=0.0,
                                     stddevs=1.0,
                                     minvals=-2.0,
                                     maxvals=2.0,
                                     dtype=dtypes.float32,
                                     seed=None,
                                     name=None):
    """Outputs random values from a truncated normal distribution.
  
    The generated values follow a normal distribution with specified mean and
    standard deviation, except that values whose magnitude is more than 2 standard
    deviations from the mean are dropped and re-picked.
  
    Args:
      shape: A 1-D integer Tensor or Python array. The shape of the output tensor.
      means: A 0-D Tensor or Python value of type `dtype`. The mean of the
        truncated normal distribution.
      stddevs: A 0-D Tensor or Python value of type `dtype`. The standard
        deviation of the truncated normal distribution.
      minvals: A 0-D Tensor or Python value of type `dtype`. The minimum value of
        the truncated normal distribution.
      maxvals: A 0-D Tensor or Python value of type `dtype`. The maximum value of
        the truncated normal distribution.
      dtype: The type of the output.
      seed: A Python integer. Used to create a random seed for the distribution.
        See
        `tf.set_random_seed`
        for behavior.
      name: A name for the operation (optional).
  
    Returns:
      A tensor of the specified shape filled with random truncated normal values.
    """
    with ops.name_scope(name, "parameterized_truncated_normal",
                        [shape, means, stddevs, minvals, maxvals]) as name:
      shape_tensor = _ShapeTensor(shape)
      means_tensor = ops.convert_to_tensor(means, dtype=dtype, name="means")
      stddevs_tensor = ops.convert_to_tensor(stddevs, dtype=dtype, name="stddevs")
      minvals_tensor = ops.convert_to_tensor(minvals, dtype=dtype, name="minvals")
      maxvals_tensor = ops.convert_to_tensor(maxvals, dtype=dtype, name="maxvals")
      seed1, seed2 = random_seed.get_seed(seed)
      rnd = gen_random_ops.parameterized_truncated_normal(
          shape_tensor,
          means_tensor,
          stddevs_tensor,
          minvals_tensor,
          maxvals_tensor,
          seed=seed1,
          seed2=seed2)
      return rnd
  
  
  @tf_export("random.truncated_normal", "truncated_normal")
  def truncated_normal(shape,
                       mean=0.0,
                       stddev=1.0,
                       dtype=dtypes.float32,
                       seed=None,
                       name=None):
    """Outputs random values from a truncated normal distribution.
  
    The generated values follow a normal distribution with specified mean and
    standard deviation, except that values whose magnitude is more than 2 standard
    deviations from the mean are dropped and re-picked.
  
    Args:
      shape: A 1-D integer Tensor or Python array. The shape of the output tensor.
      mean: A 0-D Tensor or Python value of type `dtype`. The mean of the
        truncated normal distribution.
      stddev: A 0-D Tensor or Python value of type `dtype`. The standard deviation
        of the normal distribution, before truncation.
      dtype: The type of the output.
      seed: A Python integer. Used to create a random seed for the distribution.
        See
        `tf.set_random_seed`
        for behavior.
      name: A name for the operation (optional).
  
    Returns:
      A tensor of the specified shape filled with random truncated normal values.
    """
    with ops.name_scope(name, "truncated_normal", [shape, mean, stddev]) as name:
      shape_tensor = _ShapeTensor(shape)
      mean_tensor = ops.convert_to_tensor(mean, dtype=dtype, name="mean")
      stddev_tensor = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev")
      seed1, seed2 = random_seed.get_seed(seed)
      rnd = gen_random_ops.truncated_normal(
          shape_tensor, dtype, seed=seed1, seed2=seed2)
      mul = rnd * stddev_tensor
      value = math_ops.add(mul, mean_tensor, name=name)
      return value
  
  
  ops.NotDifferentiable("ParameterizedTruncatedNormal")
  ops.NotDifferentiable("TruncatedNormal")
  
  
  @tf_export("random.uniform", "random_uniform")
  def random_uniform(shape,
                     minval=0,
                     maxval=None,
                     dtype=dtypes.float32,
                     seed=None,
                     name=None):
    """Outputs random values from a uniform distribution.
  
    The generated values follow a uniform distribution in the range
    `[minval, maxval)`. The lower bound `minval` is included in the range, while
    the upper bound `maxval` is excluded.
  
    For floats, the default range is `[0, 1)`.  For ints, at least `maxval` must
    be specified explicitly.
  
    In the integer case, the random integers are slightly biased unless
    `maxval - minval` is an exact power of two.  The bias is small for values of
    `maxval - minval` significantly smaller than the range of the output (either
    `2**32` or `2**64`).
  
    Args:
      shape: A 1-D integer Tensor or Python array. The shape of the output tensor.
      minval: A 0-D Tensor or Python value of type `dtype`. The lower bound on the
        range of random values to generate.  Defaults to 0.
      maxval: A 0-D Tensor or Python value of type `dtype`. The upper bound on
        the range of random values to generate.  Defaults to 1 if `dtype` is
        floating point.
      dtype: The type of the output: `float16`, `float32`, `float64`, `int32`,
        or `int64`.
      seed: A Python integer. Used to create a random seed for the distribution.
        See `tf.set_random_seed`
        for behavior.
      name: A name for the operation (optional).
  
    Returns:
      A tensor of the specified shape filled with random uniform values.
  
    Raises:
      ValueError: If `dtype` is integral and `maxval` is not specified.
    """
    dtype = dtypes.as_dtype(dtype)
    if dtype not in (dtypes.float16, dtypes.bfloat16, dtypes.float32,
                     dtypes.float64, dtypes.int32, dtypes.int64):
      raise ValueError("Invalid dtype %r" % dtype)
    if maxval is None:
      if dtype.is_integer:
        raise ValueError("Must specify maxval for integer dtype %r" % dtype)
      maxval = 1
    with ops.name_scope(name, "random_uniform", [shape, minval, maxval]) as name:
      shape = _ShapeTensor(shape)
      minval = ops.convert_to_tensor(minval, dtype=dtype, name="min")
      maxval = ops.convert_to_tensor(maxval, dtype=dtype, name="max")
      seed1, seed2 = random_seed.get_seed(seed)
      if dtype.is_integer:
        return gen_random_ops.random_uniform_int(
            shape, minval, maxval, seed=seed1, seed2=seed2, name=name)
      else:
        rnd = gen_random_ops.random_uniform(shape, dtype, seed=seed1, seed2=seed2)
        return math_ops.add(rnd * (maxval - minval), minval, name=name)
  
  
  ops.NotDifferentiable("RandomUniform")
  
  
  @tf_export("random.shuffle", "random_shuffle")
  def random_shuffle(value, seed=None, name=None):
    """Randomly shuffles a tensor along its first dimension.
  
    The tensor is shuffled along dimension 0, such that each `value[j]` is mapped
    to one and only one `output[i]`. For example, a mapping that might occur for a
    3x2 tensor is:
  
    ```python
    [[1, 2],       [[5, 6],
     [3, 4],  ==>   [1, 2],
     [5, 6]]        [3, 4]]
    ```
  
    Args:
      value: A Tensor to be shuffled.
      seed: A Python integer. Used to create a random seed for the distribution.
        See
        `tf.set_random_seed`
        for behavior.
      name: A name for the operation (optional).
  
    Returns:
      A tensor of same shape and type as `value`, shuffled along its first
      dimension.
    """
    seed1, seed2 = random_seed.get_seed(seed)
    return gen_random_ops.random_shuffle(
        value, seed=seed1, seed2=seed2, name=name)
  
  
  @tf_export("image.random_crop", "random_crop")
  def random_crop(value, size, seed=None, name=None):
    """Randomly crops a tensor to a given size.
  
    Slices a shape `size` portion out of `value` at a uniformly chosen offset.
    Requires `value.shape >= size`.
  
    If a dimension should not be cropped, pass the full size of that dimension.
    For example, RGB images can be cropped with
    `size = [crop_height, crop_width, 3]`.
  
    Args:
      value: Input tensor to crop.
      size: 1-D tensor with size the rank of `value`.
      seed: Python integer. Used to create a random seed. See
        `tf.set_random_seed`
        for behavior.
      name: A name for this operation (optional).
  
    Returns:
      A cropped tensor of the same rank as `value` and shape `size`.
    """
    # TODO(shlens): Implement edge case to guarantee output size dimensions.
    # If size > value.shape, zero pad the result so that it always has shape
    # exactly size.
    with ops.name_scope(name, "random_crop", [value, size]) as name:
      value = ops.convert_to_tensor(value, name="value")
      size = ops.convert_to_tensor(size, dtype=dtypes.int32, name="size")
      shape = array_ops.shape(value)
      check = control_flow_ops.Assert(
          math_ops.reduce_all(shape >= size),
          ["Need value.shape >= size, got ", shape, size],
          summarize=1000)
      shape = control_flow_ops.with_dependencies([check], shape)
      limit = shape - size + 1
      offset = random_uniform(
          array_ops.shape(shape),
          dtype=size.dtype,
          maxval=size.dtype.max,
          seed=seed) % limit
      return array_ops.slice(value, offset, size, name=name)
  
  
  @tf_export("random.multinomial", "multinomial")
  def multinomial(logits, num_samples, seed=None, name=None, output_dtype=None):
    """Draws samples from a multinomial distribution.
  
    Example:
  
    ```python
    # samples has shape [1, 5], where each value is either 0 or 1 with equal
    # probability.
    samples = tf.multinomial(tf.log([[10., 10.]]), 5)
    ```
  
    Args:
      logits: 2-D Tensor with shape `[batch_size, num_classes]`.  Each slice
        `[i, :]` represents the unnormalized log-probabilities for all classes.
      num_samples: 0-D.  Number of independent samples to draw for each row slice.
      seed: A Python integer. Used to create a random seed for the distribution.
        See
        `tf.set_random_seed`
        for behavior.
      name: Optional name for the operation.
      output_dtype: integer type to use for the output. Defaults to int64.
  
    Returns:
      The drawn samples of shape `[batch_size, num_samples]`.
    """
    with ops.name_scope(name, "multinomial", [logits]):
      logits = ops.convert_to_tensor(logits, name="logits")
      seed1, seed2 = random_seed.get_seed(seed)
      return gen_random_ops.multinomial(
          logits, num_samples, seed=seed1, seed2=seed2, output_dtype=output_dtype)
  
  
  ops.NotDifferentiable("Multinomial")
  
  
  @tf_export("random.gamma", "random_gamma")
  @deprecation.deprecated_endpoints("random_gamma")
  def random_gamma(shape,
                   alpha,
                   beta=None,
                   dtype=dtypes.float32,
                   seed=None,
                   name=None):
    """Draws `shape` samples from each of the given Gamma distribution(s).
  
    `alpha` is the shape parameter describing the distribution(s), and `beta` is
    the inverse scale parameter(s).
  
    Note: Because internal calculations are done using `float64` and casting has
    `floor` semantics, we must manually map zero outcomes to the smallest
    possible positive floating-point value, i.e., `np.finfo(dtype).tiny`.  This
    means that `np.finfo(dtype).tiny` occurs more frequently than it otherwise
    should.  This bias can only happen for small values of `alpha`, i.e.,
    `alpha << 1` or large values of `beta`, i.e., `beta >> 1`.
  
    The samples are differentiable w.r.t. alpha and beta.
    The derivatives are computed using the approach described in the paper
  
    [Michael Figurnov, Shakir Mohamed, Andriy Mnih.
    Implicit Reparameterization Gradients, 2018](https://arxiv.org/abs/1805.08498)
  
    Example:
  
    ```python
    samples = tf.random_gamma([10], [0.5, 1.5])
    # samples has shape [10, 2], where each slice [:, 0] and [:, 1] represents
    # the samples drawn from each distribution
  
    samples = tf.random_gamma([7, 5], [0.5, 1.5])
    # samples has shape [7, 5, 2], where each slice [:, :, 0] and [:, :, 1]
    # represents the 7x5 samples drawn from each of the two distributions
  
    alpha = tf.constant([[1.],[3.],[5.]])
    beta = tf.constant([[3., 4.]])
    samples = tf.random_gamma([30], alpha=alpha, beta=beta)
    # samples has shape [30, 3, 2], with 30 samples each of 3x2 distributions.
  
    loss = tf.reduce_mean(tf.square(samples))
    dloss_dalpha, dloss_dbeta = tf.gradients(loss, [alpha, beta])
    # unbiased stochastic derivatives of the loss function
    alpha.shape == dloss_dalpha.shape  # True
    beta.shape == dloss_dbeta.shape  # True
    ```
  
    Args:
      shape: A 1-D integer Tensor or Python array. The shape of the output samples
        to be drawn per alpha/beta-parameterized distribution.
      alpha: A Tensor or Python value or N-D array of type `dtype`. `alpha`
        provides the shape parameter(s) describing the gamma distribution(s) to
        sample. Must be broadcastable with `beta`.
      beta: A Tensor or Python value or N-D array of type `dtype`. Defaults to 1.
        `beta` provides the inverse scale parameter(s) of the gamma
        distribution(s) to sample. Must be broadcastable with `alpha`.
      dtype: The type of alpha, beta, and the output: `float16`, `float32`, or
        `float64`.
      seed: A Python integer. Used to create a random seed for the distributions.
        See
        `tf.set_random_seed`
        for behavior.
      name: Optional name for the operation.
  
    Returns:
      samples: a `Tensor` of shape
        `tf.concat([shape, tf.shape(alpha + beta)], axis=0)` with values of type
        `dtype`.
    """
    with ops.name_scope(name, "random_gamma", [shape, alpha, beta]):
      shape = ops.convert_to_tensor(shape, name="shape", dtype=dtypes.int32)
      alpha = ops.convert_to_tensor(alpha, name="alpha", dtype=dtype)
      beta = ops.convert_to_tensor(
          beta if beta is not None else 1, name="beta", dtype=dtype)
      alpha_broadcast = alpha + array_ops.zeros_like(beta)
      seed1, seed2 = random_seed.get_seed(seed)
      return math_ops.maximum(
          np.finfo(dtype.as_numpy_dtype).tiny,
          gen_random_ops.random_gamma(
              shape, alpha_broadcast, seed=seed1, seed2=seed2) / beta)
  
  
  @tf_export("random.poisson", "random_poisson")
  @deprecation.deprecated_endpoints("random_poisson")
  def random_poisson(lam, shape, dtype=dtypes.float32, seed=None, name=None):
    """Draws `shape` samples from each of the given Poisson distribution(s).
  
    `lam` is the rate parameter describing the distribution(s).
  
    Example:
  
    ```python
    samples = tf.random_poisson([0.5, 1.5], [10])
    # samples has shape [10, 2], where each slice [:, 0] and [:, 1] represents
    # the samples drawn from each distribution
  
    samples = tf.random_poisson([12.2, 3.3], [7, 5])
    # samples has shape [7, 5, 2], where each slice [:, :, 0] and [:, :, 1]
    # represents the 7x5 samples drawn from each of the two distributions
    ```
  
    Args:
      lam: A Tensor or Python value or N-D array of type `dtype`.
        `lam` provides the rate parameter(s) describing the poisson
        distribution(s) to sample.
      shape: A 1-D integer Tensor or Python array. The shape of the output samples
        to be drawn per "rate"-parameterized distribution.
      dtype: The type of the output: `float16`, `float32`, `float64`, `int32` or
        `int64`.
      seed: A Python integer. Used to create a random seed for the distributions.
        See
        `tf.set_random_seed`
        for behavior.
      name: Optional name for the operation.
  
    Returns:
      samples: a `Tensor` of shape `tf.concat([shape, tf.shape(lam)], axis=0)`
        with values of type `dtype`.
    """
    with ops.name_scope(name, "random_poisson", [lam, shape]):
      shape = ops.convert_to_tensor(shape, name="shape", dtype=dtypes.int32)
      seed1, seed2 = random_seed.get_seed(seed)
      return gen_random_ops.random_poisson_v2(
          shape, lam, dtype=dtype, seed=seed1, seed2=seed2)
   # Copyright 2015 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  # pylint: disable=invalid-name
  """Save and restore variables."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import collections
  import os.path
  import re
  import time
  
  from google.protobuf import text_format
  
  from tensorflow.core.protobuf import saver_pb2
  from tensorflow.python.eager import context
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.lib.io import file_io
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.platform import tf_logging as logging
  from tensorflow.python.training import training_util
  from tensorflow.python.training.checkpoint_state_pb2 import CheckpointState
  from tensorflow.python.util import compat
  from tensorflow.python.util.tf_export import tf_export
  
  
  def _GetCheckpointFilename(save_dir, latest_filename):
    """Returns a filename for storing the CheckpointState.
  
    Args:
      save_dir: The directory for saving and restoring checkpoints.
      latest_filename: Name of the file in 'save_dir' that is used
        to store the CheckpointState.
  
    Returns:
      The path of the file that contains the CheckpointState proto.
    """
    if latest_filename is None:
      latest_filename = "checkpoint"
    return os.path.join(save_dir, latest_filename)
  
  
  @tf_export("train.generate_checkpoint_state_proto")
  def generate_checkpoint_state_proto(save_dir,
                                      model_checkpoint_path,
                                      all_model_checkpoint_paths=None,
                                      all_model_checkpoint_timestamps=None,
                                      last_preserved_timestamp=None):
    """Generates a checkpoint state proto.
  
    Args:
      save_dir: Directory where the model was saved.
      model_checkpoint_path: The checkpoint file.
      all_model_checkpoint_paths: List of strings.  Paths to all not-yet-deleted
        checkpoints, sorted from oldest to newest.  If this is a non-empty list,
        the last element must be equal to model_checkpoint_path.  These paths
        are also saved in the CheckpointState proto.
      all_model_checkpoint_timestamps: A list of floats, indicating the number of
        seconds since the Epoch when each checkpoint was generated.
      last_preserved_timestamp: A float, indicating the number of seconds since
        the Epoch when the last preserved checkpoint was written, e.g. due to a
        `keep_checkpoint_every_n_hours` parameter (see
        `tf.contrib.checkpoint.CheckpointManager` for an implementation).
    Returns:
      CheckpointState proto with model_checkpoint_path and
      all_model_checkpoint_paths updated to either absolute paths or
      relative paths to the current save_dir.
  
    Raises:
      ValueError: If `all_model_checkpoint_timestamps` was provided but its length
        does not match `all_model_checkpoint_paths`.
    """
    if all_model_checkpoint_paths is None:
      all_model_checkpoint_paths = []
  
    if (not all_model_checkpoint_paths or
        all_model_checkpoint_paths[-1] != model_checkpoint_path):
      logging.info("%s is not in all_model_checkpoint_paths. Manually adding it.",
                   model_checkpoint_path)
      all_model_checkpoint_paths.append(model_checkpoint_path)
  
    if (all_model_checkpoint_timestamps
        and (len(all_model_checkpoint_timestamps)
             != len(all_model_checkpoint_paths))):
      raise ValueError(
          ("Checkpoint timestamps, if provided, must match checkpoint paths (got "
           "paths %s and timestamps %s)")
          % (all_model_checkpoint_paths, all_model_checkpoint_timestamps))
  
    # Relative paths need to be rewritten to be relative to the "save_dir"
    # if model_checkpoint_path already contains "save_dir".
    if not os.path.isabs(save_dir):
      if not os.path.isabs(model_checkpoint_path):
        model_checkpoint_path = os.path.relpath(model_checkpoint_path, save_dir)
      for i in range(len(all_model_checkpoint_paths)):
        p = all_model_checkpoint_paths[i]
        if not os.path.isabs(p):
          all_model_checkpoint_paths[i] = os.path.relpath(p, save_dir)
  
    coord_checkpoint_proto = CheckpointState(
        model_checkpoint_path=model_checkpoint_path,
        all_model_checkpoint_paths=all_model_checkpoint_paths,
        all_model_checkpoint_timestamps=all_model_checkpoint_timestamps,
        last_preserved_timestamp=last_preserved_timestamp)
  
    return coord_checkpoint_proto
  
  
  @tf_export("train.update_checkpoint_state")
  def update_checkpoint_state(save_dir,
                              model_checkpoint_path,
                              all_model_checkpoint_paths=None,
                              latest_filename=None,
                              all_model_checkpoint_timestamps=None,
                              last_preserved_timestamp=None):
    """Updates the content of the 'checkpoint' file.
  
    This updates the checkpoint file containing a CheckpointState
    proto.
  
    Args:
      save_dir: Directory where the model was saved.
      model_checkpoint_path: The checkpoint file.
      all_model_checkpoint_paths: List of strings.  Paths to all not-yet-deleted
        checkpoints, sorted from oldest to newest.  If this is a non-empty list,
        the last element must be equal to model_checkpoint_path.  These paths
        are also saved in the CheckpointState proto.
      latest_filename: Optional name of the checkpoint file.  Default to
        'checkpoint'.
      all_model_checkpoint_timestamps: Optional list of timestamps (floats,
        seconds since the Epoch) indicating when the checkpoints in
        `all_model_checkpoint_paths` were created.
      last_preserved_timestamp: A float, indicating the number of seconds since
        the Epoch when the last preserved checkpoint was written, e.g. due to a
        `keep_checkpoint_every_n_hours` parameter (see
        `tf.contrib.checkpoint.CheckpointManager` for an implementation).
    Raises:
      RuntimeError: If any of the model checkpoint paths conflict with the file
        containing CheckpointSate.
    """
    update_checkpoint_state_internal(
        save_dir=save_dir,
        model_checkpoint_path=model_checkpoint_path,
        all_model_checkpoint_paths=all_model_checkpoint_paths,
        latest_filename=latest_filename,
        save_relative_paths=False,
        all_model_checkpoint_timestamps=all_model_checkpoint_timestamps,
        last_preserved_timestamp=last_preserved_timestamp)
  
  
  def update_checkpoint_state_internal(save_dir,
                                       model_checkpoint_path,
                                       all_model_checkpoint_paths=None,
                                       latest_filename=None,
                                       save_relative_paths=False,
                                       all_model_checkpoint_timestamps=None,
                                       last_preserved_timestamp=None):
    """Updates the content of the 'checkpoint' file.
  
    This updates the checkpoint file containing a CheckpointState
    proto.
  
    Args:
      save_dir: Directory where the model was saved.
      model_checkpoint_path: The checkpoint file.
      all_model_checkpoint_paths: List of strings.  Paths to all not-yet-deleted
        checkpoints, sorted from oldest to newest.  If this is a non-empty list,
        the last element must be equal to model_checkpoint_path.  These paths
        are also saved in the CheckpointState proto.
      latest_filename: Optional name of the checkpoint file.  Default to
        'checkpoint'.
      save_relative_paths: If `True`, will write relative paths to the checkpoint
        state file.
      all_model_checkpoint_timestamps: Optional list of timestamps (floats,
        seconds since the Epoch) indicating when the checkpoints in
        `all_model_checkpoint_paths` were created.
      last_preserved_timestamp: A float, indicating the number of seconds since
        the Epoch when the last preserved checkpoint was written, e.g. due to a
        `keep_checkpoint_every_n_hours` parameter (see
        `tf.contrib.checkpoint.CheckpointManager` for an implementation).
  
    Raises:
      RuntimeError: If any of the model checkpoint paths conflict with the file
        containing CheckpointSate.
    """
    # Writes the "checkpoint" file for the coordinator for later restoration.
    coord_checkpoint_filename = _GetCheckpointFilename(save_dir, latest_filename)
    if save_relative_paths:
      if os.path.isabs(model_checkpoint_path):
        rel_model_checkpoint_path = os.path.relpath(
            model_checkpoint_path, save_dir)
      else:
        rel_model_checkpoint_path = model_checkpoint_path
      rel_all_model_checkpoint_paths = []
      for p in all_model_checkpoint_paths:
        if os.path.isabs(p):
          rel_all_model_checkpoint_paths.append(os.path.relpath(p, save_dir))
        else:
          rel_all_model_checkpoint_paths.append(p)
      ckpt = generate_checkpoint_state_proto(
          save_dir,
          rel_model_checkpoint_path,
          all_model_checkpoint_paths=rel_all_model_checkpoint_paths,
          all_model_checkpoint_timestamps=all_model_checkpoint_timestamps,
          last_preserved_timestamp=last_preserved_timestamp)
    else:
      ckpt = generate_checkpoint_state_proto(
          save_dir,
          model_checkpoint_path,
          all_model_checkpoint_paths=all_model_checkpoint_paths,
          all_model_checkpoint_timestamps=all_model_checkpoint_timestamps,
          last_preserved_timestamp=last_preserved_timestamp)
  
    if coord_checkpoint_filename == ckpt.model_checkpoint_path:
      raise RuntimeError("Save path '%s' conflicts with path used for "
                         "checkpoint state.  Please use a different save path." %
                         model_checkpoint_path)
  
    # Preventing potential read/write race condition by *atomically* writing to a
    # file.
    file_io.atomic_write_string_to_file(coord_checkpoint_filename,
                                        text_format.MessageToString(ckpt))
  
  
  @tf_export("train.get_checkpoint_state")
  def get_checkpoint_state(checkpoint_dir, latest_filename=None):
    """Returns CheckpointState proto from the "checkpoint" file.
  
    If the "checkpoint" file contains a valid CheckpointState
    proto, returns it.
  
    Args:
      checkpoint_dir: The directory of checkpoints.
      latest_filename: Optional name of the checkpoint file.  Default to
        'checkpoint'.
  
    Returns:
      A CheckpointState if the state was available, None
      otherwise.
  
    Raises:
      ValueError: if the checkpoint read doesn't have model_checkpoint_path set.
    """
    ckpt = None
    coord_checkpoint_filename = _GetCheckpointFilename(checkpoint_dir,
                                                       latest_filename)
    f = None
    try:
      # Check that the file exists before opening it to avoid
      # many lines of errors from colossus in the logs.
      if file_io.file_exists(coord_checkpoint_filename):
        file_content = file_io.read_file_to_string(
            coord_checkpoint_filename)
        ckpt = CheckpointState()
        text_format.Merge(file_content, ckpt)
        if not ckpt.model_checkpoint_path:
          raise ValueError("Invalid checkpoint state loaded from "
                           + checkpoint_dir)
        # For relative model_checkpoint_path and all_model_checkpoint_paths,
        # prepend checkpoint_dir.
        if not os.path.isabs(ckpt.model_checkpoint_path):
          ckpt.model_checkpoint_path = os.path.join(checkpoint_dir,
                                                    ckpt.model_checkpoint_path)
        for i in range(len(ckpt.all_model_checkpoint_paths)):
          p = ckpt.all_model_checkpoint_paths[i]
          if not os.path.isabs(p):
            ckpt.all_model_checkpoint_paths[i] = os.path.join(checkpoint_dir, p)
    except errors.OpError as e:
      # It's ok if the file cannot be read
      logging.warning("%s: %s", type(e).__name__, e)
      logging.warning("%s: Checkpoint ignored", coord_checkpoint_filename)
      return None
    except text_format.ParseError as e:
      logging.warning("%s: %s", type(e).__name__, e)
      logging.warning("%s: Checkpoint ignored", coord_checkpoint_filename)
      return None
    finally:
      if f:
        f.close()
    return ckpt
  
  
  def _prefix_to_checkpoint_path(prefix, format_version):
    """Returns the pathname of a checkpoint file, given the checkpoint prefix.
  
    For V1 checkpoint, simply returns the prefix itself (the data file).  For V2,
    returns the pathname to the index file.
  
    Args:
      prefix: a string, the prefix of a checkpoint.
      format_version: the checkpoint format version that corresponds to the
        prefix.
    Returns:
      The pathname of a checkpoint file, taking into account the checkpoint
        format version.
    """
    if format_version == saver_pb2.SaverDef.V2:
      return prefix + ".index"  # The index file identifies a checkpoint.
    return prefix  # Just the data file.
  
  
  @tf_export("train.latest_checkpoint")
  def latest_checkpoint(checkpoint_dir, latest_filename=None):
    """Finds the filename of latest saved checkpoint file.
  
    Args:
      checkpoint_dir: Directory where the variables were saved.
      latest_filename: Optional name for the protocol buffer file that
        contains the list of most recent checkpoint filenames.
        See the corresponding argument to `Saver.save()`.
  
    Returns:
      The full path to the latest checkpoint or `None` if no checkpoint was found.
    """
    # Pick the latest checkpoint based on checkpoint state.
    ckpt = get_checkpoint_state(checkpoint_dir, latest_filename)
    if ckpt and ckpt.model_checkpoint_path:
      # Look for either a V2 path or a V1 path, with priority for V2.
      v2_path = _prefix_to_checkpoint_path(ckpt.model_checkpoint_path,
                                           saver_pb2.SaverDef.V2)
      v1_path = _prefix_to_checkpoint_path(ckpt.model_checkpoint_path,
                                           saver_pb2.SaverDef.V1)
      if file_io.get_matching_files(v2_path) or file_io.get_matching_files(
          v1_path):
        return ckpt.model_checkpoint_path
      else:
        logging.error("Couldn't match files for checkpoint %s",
                      ckpt.model_checkpoint_path)
    return None
  
  
  @tf_export("train.checkpoint_exists")
  def checkpoint_exists(checkpoint_prefix):
    """Checks whether a V1 or V2 checkpoint exists with the specified prefix.
  
    This is the recommended way to check if a checkpoint exists, since it takes
    into account the naming difference between V1 and V2 formats.
  
    Args:
      checkpoint_prefix: the prefix of a V1 or V2 checkpoint, with V2 taking
        priority.  Typically the result of `Saver.save()` or that of
        `tf.train.latest_checkpoint()`, regardless of sharded/non-sharded or
        V1/V2.
    Returns:
      A bool, true iff a checkpoint referred to by `checkpoint_prefix` exists.
    """
    pathname = _prefix_to_checkpoint_path(checkpoint_prefix,
                                          saver_pb2.SaverDef.V2)
    if file_io.get_matching_files(pathname):
      return True
    elif file_io.get_matching_files(checkpoint_prefix):
      return True
    else:
      return False
  
  
  @tf_export("train.get_checkpoint_mtimes")
  def get_checkpoint_mtimes(checkpoint_prefixes):
    """Returns the mtimes (modification timestamps) of the checkpoints.
  
    Globs for the checkpoints pointed to by `checkpoint_prefixes`.  If the files
    exist, collect their mtime.  Both V2 and V1 checkpoints are considered, in
    that priority.
  
    This is the recommended way to get the mtimes, since it takes into account
    the naming difference between V1 and V2 formats.
  
    Args:
      checkpoint_prefixes: a list of checkpoint paths, typically the results of
        `Saver.save()` or those of `tf.train.latest_checkpoint()`, regardless of
        sharded/non-sharded or V1/V2.
    Returns:
      A list of mtimes (in microseconds) of the found checkpoints.
    """
    mtimes = []
  
    def match_maybe_append(pathname):
      fnames = file_io.get_matching_files(pathname)
      if fnames:
        mtimes.append(file_io.stat(fnames[0]).mtime_nsec / 1e9)
        return True
      return False
  
    for checkpoint_prefix in checkpoint_prefixes:
      # Tries V2's metadata file first.
      pathname = _prefix_to_checkpoint_path(checkpoint_prefix,
                                            saver_pb2.SaverDef.V2)
      if match_maybe_append(pathname):
        continue
      # Otherwise, tries V1, where the prefix is the complete pathname.
      match_maybe_append(checkpoint_prefix)
  
    return mtimes
  
  
  @tf_export("train.remove_checkpoint")
  def remove_checkpoint(checkpoint_prefix,
                        checkpoint_format_version=saver_pb2.SaverDef.V2,
                        meta_graph_suffix="meta"):
    """Removes a checkpoint given by `checkpoint_prefix`.
  
    Args:
      checkpoint_prefix: The prefix of a V1 or V2 checkpoint. Typically the result
        of `Saver.save()` or that of `tf.train.latest_checkpoint()`, regardless of
        sharded/non-sharded or V1/V2.
      checkpoint_format_version: `SaverDef.CheckpointFormatVersion`, defaults to
        `SaverDef.V2`.
      meta_graph_suffix: Suffix for `MetaGraphDef` file. Defaults to 'meta'.
    """
    _delete_file_if_exists(
        meta_graph_filename(checkpoint_prefix, meta_graph_suffix))
    if checkpoint_format_version == saver_pb2.SaverDef.V2:
      # V2 has a metadata file and some data files.
      _delete_file_if_exists(checkpoint_prefix + ".index")
      _delete_file_if_exists(checkpoint_prefix + ".data-?????-of-?????")
    else:
      # V1, Legacy.  Exact match on the data file.
      _delete_file_if_exists(checkpoint_prefix)
  
  
  def _delete_file_if_exists(filespec):
    """Deletes files matching `filespec`."""
    for pathname in file_io.get_matching_files(filespec):
      file_io.delete_file(pathname)
  
  
  def meta_graph_filename(checkpoint_filename, meta_graph_suffix="meta"):
    """Returns the meta graph filename.
  
    Args:
      checkpoint_filename: Name of the checkpoint file.
      meta_graph_suffix: Suffix for `MetaGraphDef` file. Defaults to 'meta'.
  
    Returns:
      MetaGraph file name.
    """
    # If the checkpoint_filename is sharded, the checkpoint_filename could
    # be of format model.ckpt-step#-?????-of-shard#. For example,
    # model.ckpt-123456-?????-of-00005, or model.ckpt-123456-00001-of-00002.
    basename = re.sub(r"-[\d\?]+-of-\d+$", "", checkpoint_filename)
    suffixed_filename = ".".join([basename, meta_graph_suffix])
    return suffixed_filename
  
  
  # TODO(allenl): Allow tf.keras.Model instances in the constructor directly?
  class CheckpointManager(object):
    """Deletes old checkpoints.
  
    Example usage:
    ```python
    import tensorflow as tf
    checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
    manager = tf.contrib.checkpoint.CheckpointManager(
        checkpoint, directory="/tmp/model", max_to_keep=5)
    status = checkpoint.restore(manager.latest_checkpoint)
    while True:
      # train
      manager.save()
    ```
  
    `CheckpointManager` preserves its own state across instantiations (see the
    `__init__` documentation for details). Only one should be active in a
    particular directory at a time.
    """
  
    def __init__(self, checkpoint, directory,
                 max_to_keep, keep_checkpoint_every_n_hours=None):
      """Configure a `CheckpointManager` for use in `directory`.
  
      If a `CheckpointManager` was previously used in `directory`, its
      state will be restored. This includes the list of managed checkpoints and
      the timestamp bookkeeping necessary to support
      `keep_checkpoint_every_n_hours`. The behavior of the new `CheckpointManager`
      will be the same as the previous `CheckpointManager`, including cleaning up
      existing checkpoints if appropriate.
  
      Checkpoints are only considered for deletion just after a new checkpoint has
      been added. At that point, `max_to_keep` checkpoints will remain in an
      "active set". Once a checkpoint is preserved by
      `keep_checkpoint_every_n_hours` it will not be deleted by this
      `CheckpointManager` or any future `CheckpointManager` instantiated in
      `directory` (regardless of the new setting of
      `keep_checkpoint_every_n_hours`). The `max_to_keep` checkpoints in the
      active set may be deleted by this `CheckpointManager` or a future
      `CheckpointManager` instantiated in `directory` (subject to its
      `max_to_keep` and `keep_checkpoint_every_n_hours` settings).
  
      Args:
        checkpoint: The `tf.train.Checkpoint` instance to save and manage
          checkpoints for.
        directory: The path to a directory in which to write checkpoints. A
          special file named "checkpoint" is also written to this directory (in a
          human-readable text format) which contains the state of the
          `CheckpointManager`.
        max_to_keep: An integer, the number of checkpoints to keep. Unless
          preserved by `keep_checkpoint_every_n_hours`, checkpoints will be
          deleted from the active set, oldest first, until only `max_to_keep`
          checkpoints remain. If `None`, no checkpoints are deleted and everything
          stays in the active set. Note that `max_to_keep=None` will keep all
          checkpoint paths in memory and in the checkpoint state protocol buffer
          on disk.
        keep_checkpoint_every_n_hours: Upon removal from the active set, a
          checkpoint will be preserved if it has been at least
          `keep_checkpoint_every_n_hours` since the last preserved checkpoint. The
          default setting of `None` does not preserve any checkpoints in this way.
  
      Raises:
        ValueError: If `max_to_keep` is not a positive integer.
      """
      self._checkpoint = checkpoint
      self._save_counter_assign = None
      if max_to_keep is not None and max_to_keep <= 0:
        raise ValueError(
            ("Expected a positive integer or `None` for `max_to_max_to_keep`, "
             "got %d.")
            % (max_to_keep,))
      self._max_to_keep = max_to_keep
      self._keep_checkpoint_every_n_hours = keep_checkpoint_every_n_hours
      self._directory = directory
      self._checkpoint_prefix = os.path.join(directory, "ckpt")
      recovered_state = get_checkpoint_state(directory)
      current_clock = time.time()
      self._maybe_delete = collections.OrderedDict()
      if recovered_state is None:
        self._latest_checkpoint = None
        # Set the clock back slightly to avoid race conditions when quckly
        # re-creating a CheckpointManager.
        self._last_preserved_timestamp = current_clock - 1.
      else:
        self._latest_checkpoint = recovered_state.model_checkpoint_path
        self._last_preserved_timestamp = recovered_state.last_preserved_timestamp
        if current_clock < self._last_preserved_timestamp:
          # Time seems to have reversed itself. In addition to this warning, we'll
          # min() saved checkpoint timestamps with the current time to ensure that
          # old checkpoints don't get deleted accidentally.
          logging.warning(
              ("time.time() returned a value %f seconds behind the last "
               "preserved checkpoint timestamp.")
              % (self._last_preserved_timestamp - current_clock,))
          self._last_preserved_timestamp = current_clock
        all_timestamps = recovered_state.all_model_checkpoint_timestamps
        all_paths = recovered_state.all_model_checkpoint_paths
        del recovered_state  # Uses modified values from now on
        if not all_timestamps:
          all_timestamps = [self._last_preserved_timestamp] * len(all_paths)
  
        for filename, timestamp in zip(all_paths, all_timestamps):
          timestamp = min(timestamp, current_clock)
          if timestamp > self._last_preserved_timestamp:
            self._maybe_delete[filename] = timestamp
  
    @property
    def latest_checkpoint(self):
      """The prefix of the most recent checkpoint in `directory`.
  
      Equivalent to `tf.train.latest_checkpoint(directory)` where `directory` is
      the constructor argument to `CheckpointManager`.
  
      Suitable for passing to `tf.train.Checkpoint.restore` to resume training.
  
      Returns:
        The checkpoint prefix. If there are no checkpoints, returns `None`.
      """
      return self._latest_checkpoint
  
    @property
    def checkpoints(self):
      """A list of managed checkpoints.
  
      Note that checkpoints saved due to `keep_checkpoint_every_n_hours` will not
      show up in this list (to avoid ever-growing filename lists).
  
      Returns:
        A list of filenames, sorted from oldest to newest.
      """
      return list(self._maybe_delete.keys())
  
    def _sweep(self):
      """Deletes or preserves managed checkpoints."""
      if not self._max_to_keep:
        # Does not update self._last_preserved_timestamp, since everything is kept
        # in the active set.
        return
      while len(self._maybe_delete) > self._max_to_keep:
        filename, timestamp = self._maybe_delete.popitem(last=False)
        # Even if we're keeping this checkpoint due to
        # keep_checkpoint_every_n_hours, we won't reference it to avoid
        # infinitely-growing CheckpointState protos.
        if (self._keep_checkpoint_every_n_hours
            and (timestamp - self._keep_checkpoint_every_n_hours * 3600.
                 >= self._last_preserved_timestamp)):
          self._last_preserved_timestamp = timestamp
          continue
        remove_checkpoint(filename)
  
    def _record_state(self):
      """Saves the `CheckpointManager`'s state in `directory`."""
      filenames, timestamps = zip(*self._maybe_delete.items())
      update_checkpoint_state_internal(
          self._directory,
          model_checkpoint_path=self.latest_checkpoint,
          all_model_checkpoint_paths=filenames,
          all_model_checkpoint_timestamps=timestamps,
          last_preserved_timestamp=self._last_preserved_timestamp,
          save_relative_paths=True)
  
    @property
    def _prefix(self):
      """A common prefix for all checkpoints saved with this manager.
  
      For example, if `directory` (a constructor argument) were `"/tmp/tf-model"`,
      `prefix` would be `"/tmp/tf-model/ckpt"` and checkpoints would generally be
      numbered `"/tmp/tf-model/ckpt-1"`, `"/tmp/tf-model/ckpt-2"`, and so on. Each
      checkpoint has several associated files
      (e.g. `"/tmp/tf-model/ckpt-2.index"`).
  
      Returns:
        A string prefix.
      """
      return self._checkpoint_prefix
  
    def save(self, session=None, checkpoint_number=None):
      """Creates a new checkpoint and manages it.
  
      Args:
        session: The session to evaluate variables in. Ignored when executing
          eagerly. If not provided when graph building, the default session is
          used.
        checkpoint_number: An optional integer, or an integer-dtype `Variable` or
          `Tensor`, used to number the checkpoint. If `None` (default),
          checkpoints are numbered using `checkpoint.save_counter`. Even if
          `checkpoint_number` is provided, `save_counter` is still incremented. A
          user-provided `checkpoint_number` is not incremented even if it is a
          `Variable`.
  
      Returns:
        The path to the new checkpoint. It is also recorded in the `checkpoints`
        and `latest_checkpoint` properies.
      """
      # Save counter logic duplicated from tf.train.Checkpoint, soon to diverge
      # slightly with a custom numbering option.
      if context.executing_eagerly():
        save_counter = self._checkpoint.save_counter
        save_counter.assign_add(1)
      else:
        if session is None:
          session = ops.get_default_session()
  
        def _initializing_creator(next_creator, **kwargs):
          """Initialize the save counter if it has been newly created."""
          v = next_creator(**kwargs)
          session.run(v.initializer)
          return v
  
        with variable_scope.variable_creator_scope(_initializing_creator):
          save_counter = self._checkpoint.save_counter
        if self._save_counter_assign is None:
          self._save_counter_assign = save_counter.assign_add(1, read_value=False)
        session.run(self._save_counter_assign)
      if checkpoint_number is None:
        checkpoint_number = save_counter
      if not isinstance(checkpoint_number, compat.integral_types):
        checkpoint_number = training_util.global_step(
            sess=session, global_step_tensor=checkpoint_number)
      prefix = "%s-%d" % (self._prefix, checkpoint_number)
      save_path = self._checkpoint.write(prefix)
      timestamp = time.time()
      # If this is an overwritten checkpoint we were previously tracking, delete
      # and reinsert it to make sure it goes to the end of the queue.
      if save_path in self._maybe_delete:
        del self._maybe_delete[save_path]
      self._maybe_delete[save_path] = timestamp
      self._latest_checkpoint = save_path
      self._sweep()
      self._record_state()
      return save_path
  
apis: tensorflow.python.ops.array_ops.reverse tensorflow.python.ops.io_ops.TextLineReader tensorflow.python.framework.tensor_shape.TensorShape tensorflow.python.ops.state_ops.is_variable_initialized tensorflow.python.framework.constant_op.constant tensorflow.python.estimator.estimator_lib.inputs.numpy_input_fn tensorflow.python.ops.io_ops.TFRecordReader tensorflow.python.training.input.batch tensorflow.python.ops.parsing_ops.parse_example tensorflow.python.ops.math_ops.reduce_all tensorflow.python.ops.parsing_ops.decode_csv tensorflow.contrib.timeseries.python.timeseries.model_utils.canonicalize_times_or_steps_from_output tensorflow.python.ops.random_ops.random_uniform tensorflow.python.ops.array_ops.identity tensorflow.python.ops.array_ops.zeros tensorflow.python.ops.math_ops.reduce_any tensorflow.python.training.input.maybe_shuffle_batch tensorflow.python.ops.array_ops.shape tensorflow.python.ops.array_ops.squeeze tensorflow.python.training.training.limit_epochs tensorflow.python.training.input.string_input_producer tensorflow.python.framework.ops.control_dependencies tensorflow.python.framework.ops.convert_to_tensor tensorflow.python.ops.array_ops.gather tensorflow.python.ops.state_ops.assign numpy.squeeze tensorflow.python.ops.tensor_array_ops.TensorArray tensorflow.python.ops.array_ops.stack numpy.array
apis: tensorflow.python.util.deprecation.deprecated_endpoints tensorflow.python.ops.gen_random_ops.parameterized_truncated_normal tensorflow.python.ops.gen_random_ops.random_uniform_int tensorflow.python.util.tf_export.tf_export tensorflow.python.ops.math_ops.add tensorflow.python.ops.array_ops.slice tensorflow.python.ops.math_ops.reduce_all tensorflow.python.ops.gen_random_ops.truncated_normal tensorflow.python.framework.dtypes.as_dtype tensorflow.python.ops.gen_random_ops.random_poisson_v2 tensorflow.python.ops.array_ops.shape tensorflow.python.ops.control_flow_ops.with_dependencies tensorflow.python.framework.ops.convert_to_tensor tensorflow.python.ops.gen_random_ops.random_standard_normal numpy.finfo tensorflow.python.ops.gen_random_ops.multinomial tensorflow.python.ops.array_ops.zeros_like tensorflow.python.framework.random_seed.get_seed tensorflow.python.framework.ops.NotDifferentiable tensorflow.python.framework.ops.name_scope tensorflow.python.ops.gen_random_ops.random_gamma tensorflow.python.ops.gen_random_ops.random_shuffle tensorflow.python.ops.gen_random_ops.random_uniform
apis: tensorflow.python.platform.tf_logging.info tensorflow.python.util.tf_export.tf_export tensorflow.python.lib.io.file_io.delete_file tensorflow.python.platform.tf_logging.error tensorflow.python.training.training_util.global_step tensorflow.python.framework.ops.get_default_session tensorflow.python.lib.io.file_io.stat tensorflow.python.ops.variable_scope.variable_creator_scope tensorflow.python.lib.io.file_io.get_matching_files tensorflow.python.eager.context.executing_eagerly tensorflow.python.lib.io.file_io.file_exists tensorflow.python.platform.tf_logging.warning tensorflow.python.lib.io.file_io.read_file_to_string tensorflow.python.training.checkpoint_state_pb2.CheckpointState