diff options
author | Manjunath Kudlur <keveman@gmail.com> | 2015-11-06 16:27:58 -0800 |
---|---|---|
committer | Manjunath Kudlur <keveman@gmail.com> | 2015-11-06 16:27:58 -0800 |
commit | f41959ccb2d9d4c722fe8fc3351401d53bcf4900 (patch) | |
tree | ef0ca22cb2a5ac4bdec9d080d8e0788a53ed496d /tensorflow/core/kernels/initializable_lookup_table.h |
TensorFlow: Initial commit of TensorFlow library.
TensorFlow is an open source software library for numerical computation
using data flow graphs.
Base CL: 107276108
Diffstat (limited to 'tensorflow/core/kernels/initializable_lookup_table.h')
-rw-r--r-- | tensorflow/core/kernels/initializable_lookup_table.h | 103 |
1 files changed, 103 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/initializable_lookup_table.h b/tensorflow/core/kernels/initializable_lookup_table.h new file mode 100644 index 0000000000..651b491457 --- /dev/null +++ b/tensorflow/core/kernels/initializable_lookup_table.h @@ -0,0 +1,103 @@ +#ifndef TENSORFLOW_KERNELS_INITIALIZABLE_LOOKUP_TABLE_H_ +#define TENSORFLOW_KERNELS_INITIALIZABLE_LOOKUP_TABLE_H_ + +#include "tensorflow/core/framework/lookup_interface.h" + +namespace tensorflow { +namespace lookup { + +// Base class for lookup tables that require initialization. +class InitializableLookupTable : public LookupInterface { + public: + class InitTableIterator; + + // Performs batch lookups, for every element in the key tensor, Find returns + // the corresponding value into the values tensor. + // If an element is not present in the table, the given default value is used. + // + // For tables that require initialization, `Find` is available once the table + // is marked as initialized. + // + // Returns the following statuses: + // - OK: when the find finishes successfully. + // - FailedPrecondition: if the table is not initialized. + // - InvalidArgument: if any of the preconditions on the lookup key or value + // fails. + // - In addition, other implementations may provide another non-OK status + // specific to their failure modes. + Status Find(const Tensor& keys, Tensor* values, + const Tensor& default_value) final; + + // Returns whether the table was initialized and is ready to serve lookups. + bool is_initialized() const { return is_initialized_; } + + // Initializes the table from the given init table iterator. + // + // Atomically, this operation prepares the table, populates it with the given + // iterator, and mark the table as initialized. + // + // Returns the following statuses: + // - OK: when the initialization was successful. + // - InvalidArgument: if any of the preconditions on the lookup key or value + // fails. + // - FailedPrecondition: if the table is already initialized and + // fail_if_initialized is set to true. + // - In addition, other implementations may provide another non-OK status + // specific to their failure modes. + Status Initialize(InitTableIterator& iter); + + // Basic iterator to initialize lookup tables. + // It yields a sequence of pairs of `keys()` and `values()` Tensors, so that + // the consumer may insert key-value pairs in batches. + // + // Then the iterator is exhausted, valid returns false and status returns + // Status::OutOfRange. + class InitTableIterator { + public: + InitTableIterator() {} + + virtual ~InitTableIterator() {} + + // Prepares the next batch of key and value tensors. + virtual void Next() = 0; + + // Returns true if keys and values point to valid tensors. + virtual bool Valid() const = 0; + + // Returns a tensor that contains the current batch of 'key' values. + virtual const Tensor& keys() const = 0; + + // Returns a tensor that contains the current batch of 'value' values. + virtual const Tensor& values() const = 0; + + // Returns an error if one has occurred, otherwire returns Status::OK. + virtual Status status() const = 0; + + // Returns the total number of elements that the iterator will produce. + virtual int64 total_size() const = 0; + + private: + TF_DISALLOW_COPY_AND_ASSIGN(InitTableIterator); + }; + + protected: + // Prepares and allocates the underlying data structure to store the given + // number of expected elements. + virtual Status DoPrepare(size_t expected_num_elements) = 0; + + // Populates the table in batches given keys and values as tensors into the + // underlying data structure. + virtual Status DoInsert(const Tensor& keys, const Tensor& values) = 0; + + // Performs the batch find operation on the underlying data structure. + virtual Status DoFind(const Tensor& keys, Tensor* values, + const Tensor& default_value) = 0; + + mutex mu_; + bool is_initialized_ = false; +}; + +} // namespace lookup +} // namespace tensorflow + +#endif // TENSORFLOW_KERNELS_INITIALIZABLE_LOOKUP_TABLE_H_ |