From 6c391166b8b6ba43d2b0151e6fb9cf14864131a2 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 9 Oct 2018 16:23:35 -0700 Subject: Add 'remove' operation to MutableHashTable and MutableDenseHashTable. PiperOrigin-RevId: 216443201 --- tensorflow/core/framework/lookup_interface.cc | 8 ++++++++ tensorflow/core/framework/lookup_interface.h | 17 +++++++++++++++++ 2 files changed, 25 insertions(+) (limited to 'tensorflow/core/framework') diff --git a/tensorflow/core/framework/lookup_interface.cc b/tensorflow/core/framework/lookup_interface.cc index bf3204ea6e..117adbf65c 100644 --- a/tensorflow/core/framework/lookup_interface.cc +++ b/tensorflow/core/framework/lookup_interface.cc @@ -71,6 +71,14 @@ Status LookupInterface::CheckKeyAndValueTensorsForImport(const Tensor& keys, return CheckKeyAndValueTensorsHelper(keys, values); } +Status LookupInterface::CheckKeyTensorForRemove(const Tensor& keys) { + if (keys.dtype() != key_dtype()) { + return errors::InvalidArgument("Key must be type ", key_dtype(), + " but got ", keys.dtype()); + } + return CheckKeyShape(keys.shape()); +} + Status LookupInterface::CheckFindArguments(const Tensor& key, const Tensor& default_value) { TF_RETURN_IF_ERROR(CheckKeyAndValueTypes(key, default_value)); diff --git a/tensorflow/core/framework/lookup_interface.h b/tensorflow/core/framework/lookup_interface.h index 0622dd06cb..d33945fd1b 100644 --- a/tensorflow/core/framework/lookup_interface.h +++ b/tensorflow/core/framework/lookup_interface.h @@ -64,6 +64,17 @@ class LookupInterface : public ResourceBase { virtual Status Insert(OpKernelContext* ctx, const Tensor& keys, const Tensor& values) = 0; + // Removes elements from the table. + // This method is only implemented in mutable tables that can be updated over + // the execution of the graph. It returns Status::NotImplemented for read-only + // tables that are initialized once before they can be looked up. + + // Returns the following statuses: + // - OK: when the remove finishes successfully. + // - InvalidArgument: if any of the preconditions on the lookup key fails. + // - Unimplemented: if the table does not support removals. + virtual Status Remove(OpKernelContext* ctx, const Tensor& keys) = 0; + // Returns the number of elements in the table. virtual size_t size() const = 0; @@ -107,6 +118,12 @@ class LookupInterface : public ResourceBase { virtual Status CheckKeyAndValueTensorsForImport(const Tensor& keys, const Tensor& values); + // Check format of the key tensor for the Remove function. + // Returns OK if all the following requirements are satisfied, otherwise it + // returns InvalidArgument: + // - DataType of the tensor keys equals to the table key_dtype + virtual Status CheckKeyTensorForRemove(const Tensor& keys); + // Check the arguments of a find operation. Returns OK if all the following // requirements are satisfied, otherwise it returns InvalidArgument: // - DataType of the tensor keys equals to the table key_dtype -- cgit v1.2.3