diff options
author | A. Unique TensorFlower <gardener@tensorflow.org> | 2018-10-09 16:23:35 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-10-09 16:39:23 -0700 |
commit | 6c391166b8b6ba43d2b0151e6fb9cf14864131a2 (patch) | |
tree | c8c3c9eadade00f1a4e6cec2024e2a15bfd0b948 /tensorflow/core/ops | |
parent | 2f5ebc0ea5e6d500ea8cd925234c569d6b32fd4e (diff) |
Add 'remove' operation to MutableHashTable and MutableDenseHashTable.
PiperOrigin-RevId: 216443201
Diffstat (limited to 'tensorflow/core/ops')
-rw-r--r-- | tensorflow/core/ops/compat/ops_history.v1.pbtxt | 20 | ||||
-rw-r--r-- | tensorflow/core/ops/lookup_ops.cc | 14 |
2 files changed, 34 insertions, 0 deletions
diff --git a/tensorflow/core/ops/compat/ops_history.v1.pbtxt b/tensorflow/core/ops/compat/ops_history.v1.pbtxt index cfb1055d3c..415e15b720 100644 --- a/tensorflow/core/ops/compat/ops_history.v1.pbtxt +++ b/tensorflow/core/ops/compat/ops_history.v1.pbtxt @@ -30321,6 +30321,22 @@ op { is_stateful: true } op { + name: "LookupTableRemoveV2" + input_arg { + name: "table_handle" + type: DT_RESOURCE + } + input_arg { + name: "keys" + type_attr: "Tin" + } + attr { + name: "Tin" + type: "type" + } + is_stateful: true +} +op { name: "LookupTableSize" input_arg { name: "table_handle" @@ -36706,6 +36722,10 @@ op { name: "empty_key" type_attr: "key_dtype" } + input_arg { + name: "deleted_key" + type_attr: "key_dtype" + } output_arg { name: "table_handle" type: DT_RESOURCE diff --git a/tensorflow/core/ops/lookup_ops.cc b/tensorflow/core/ops/lookup_ops.cc index 72a77be70d..a0987cd982 100644 --- a/tensorflow/core/ops/lookup_ops.cc +++ b/tensorflow/core/ops/lookup_ops.cc @@ -214,6 +214,19 @@ REGISTER_OP("LookupTableInsertV2") return Status::OK(); }); +REGISTER_OP("LookupTableRemoveV2") + .Input("table_handle: resource") + .Input("keys: Tin") + .Attr("Tin: type") + .SetShapeFn([](InferenceContext* c) { + ShapeHandle handle; + TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle)); + TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 1, &handle)); + + // TODO(turboale): Validate keys shape. + return Status::OK(); + }); + REGISTER_OP("LookupTableSize") .Input("table_handle: Ref(string)") .Output("size: int64") @@ -407,6 +420,7 @@ REGISTER_OP("MutableDenseHashTable") REGISTER_OP("MutableDenseHashTableV2") .Input("empty_key: key_dtype") + .Input("deleted_key: key_dtype") .Output("table_handle: resource") .Attr("container: string = ''") .Attr("shared_name: string = ''") |