1dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 3dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonLicensed under the Apache License, Version 2.0 (the "License"); 4dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leonyou may not use this file except in compliance with the License. 5dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonYou may obtain a copy of the License at 6dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 7dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon http://www.apache.org/licenses/LICENSE-2.0 8dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 9dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonUnless required by applicable law or agreed to in writing, software 10dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leondistributed under the License is distributed on an "AS IS" BASIS, 11dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonSee the License for the specific language governing permissions and 13dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leonlimitations under the License. 14dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon==============================================================================*/ 15dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 16dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon#include "tensorflow/core/framework/common_shape_fns.h" 17cf330efa207d635efecfb1b703ee65ccfbc3f98cDerek Murray#include "tensorflow/core/framework/dataset_stateful_op_whitelist.h" 18dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon#include "tensorflow/core/framework/op.h" 19dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon#include "tensorflow/core/framework/op_def_builder.h" 20dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon#include "tensorflow/core/framework/shape_inference.h" 21dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 22dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leonnamespace tensorflow { 23dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 24dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leonusing shape_inference::DimensionHandle; 25dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leonusing shape_inference::InferenceContext; 26dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leonusing shape_inference::ShapeHandle; 27dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 28dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon// -------------------------------------------------------------------------- 29dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 30dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leonnamespace { 31dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonStatus TwoElementVectorInputsAndScalarOutputs(InferenceContext* c) { 32dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 33dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon DimensionHandle unused_handle; 34dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon for (int i = 0; i < c->num_inputs(); ++i) { 35dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 1, &handle)); 36dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithValue(c->Dim(handle, 0), 2, &unused_handle)); 37dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon } 38dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon for (int i = 0; i < c->num_outputs(); ++i) { 39dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon c->set_output(i, c->Scalar()); 40dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon } 41dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 42dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon} 43dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 44dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonStatus ScalarAndTwoElementVectorInputsAndScalarOutputs(InferenceContext* c) { 45dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 46dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon DimensionHandle unused_handle; 47dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle)); 48dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon for (int i = 1; i < c->num_inputs(); ++i) { 49dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 1, &handle)); 50dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithValue(c->Dim(handle, 0), 2, &unused_handle)); 51dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon } 52dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon for (int i = 0; i < c->num_outputs(); ++i) { 53dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon c->set_output(i, c->Scalar()); 54dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon } 55dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 56dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon} 57dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 58dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonStatus TwoElementOutput(InferenceContext* c) { 59dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon c->set_output(0, c->Vector(2)); 60dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 61dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon} 62dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 63dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonStatus ScalarOutput(InferenceContext* c) { 64dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon c->set_output(0, c->Scalar()); 65dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 66dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon} 67dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon} // namespace 68dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 69dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("LookupTableFind") 70dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: Ref(string)") 71dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("keys: Tin") 72dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("default_value: Tout") 73dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("values: Tout") 74dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tin: type") 75dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tout: type") 76dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 77dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 78dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &handle)); 79dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon DimensionHandle unused_dim; 80dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithValue(c->Dim(handle, 0), 2, &unused_dim)); 81dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 82dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon // Default value must be scalar or vector. 83dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle unused; 84dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(2), 1, &unused)); 85dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon c->set_output(0, c->UnknownShape()); 86dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 873e852d462aaba446f62f76007405c0794a6087b9Anna R }); 88dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 89dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("LookupTableFindV2") 90dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: resource") 91dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("keys: Tin") 92dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("default_value: Tout") 93dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("values: Tout") 94dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tin: type") 95dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tout: type") 96dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 97dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 98dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle)); 99dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 100dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon // Default value must be scalar or vector. 101dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle unused; 102dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(2), 1, &unused)); 103dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon c->set_output(0, c->UnknownShape()); 104dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 1053e852d462aaba446f62f76007405c0794a6087b9Anna R }); 10646093eabe85895ed62a447d08b4988edf8fccb1aShivani AgrawalWHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS("LookupTableFindV2"); 10746093eabe85895ed62a447d08b4988edf8fccb1aShivani Agrawal// TODO(b/72710477): Update this. 108dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 109dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("LookupTableInsert") 110dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: Ref(string)") 111dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("keys: Tin") 112dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("values: Tout") 113dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tin: type") 114dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tout: type") 115dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 116dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 117dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &handle)); 118dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon DimensionHandle unused_dim; 119dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithValue(c->Dim(handle, 0), 2, &unused_dim)); 120dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 121dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon // TODO(ebrevdo): Validate keys and values shape. 122dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 1233e852d462aaba446f62f76007405c0794a6087b9Anna R }); 124dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 125dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("LookupTableInsertV2") 126dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: resource") 127dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("keys: Tin") 128dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("values: Tout") 129dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tin: type") 130dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tout: type") 131dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 132dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 133dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle)); 134dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 135dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon // TODO: Validate keys and values shape. 136dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 1373e852d462aaba446f62f76007405c0794a6087b9Anna R }); 138dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 139dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("LookupTableSize") 140dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: Ref(string)") 141dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("size: int64") 1423e852d462aaba446f62f76007405c0794a6087b9Anna R .SetShapeFn(TwoElementVectorInputsAndScalarOutputs); 143dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 144dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("LookupTableSizeV2") 145dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: resource") 146dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("size: int64") 1473e852d462aaba446f62f76007405c0794a6087b9Anna R .SetShapeFn(ScalarAndTwoElementVectorInputsAndScalarOutputs); 148dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 149dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("LookupTableExport") 150dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: Ref(string)") 151dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("keys: Tkeys") 152dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("values: Tvalues") 153dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tkeys: type") 154dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tvalues: type") 155dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 156dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 157dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &handle)); 158dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon DimensionHandle unused_dim; 159dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithValue(c->Dim(handle, 0), 2, &unused_dim)); 160dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 161dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle values = c->UnknownShape(); 162dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRankAtLeast(values, 1, &values)); 163dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle keys = c->Vector(c->Dim(values, 0)); 164dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon c->set_output(0, keys); 165dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon c->set_output(1, values); 166dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 1673e852d462aaba446f62f76007405c0794a6087b9Anna R }); 168dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 169dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("LookupTableExportV2") 170dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: resource") 171dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("keys: Tkeys") 172dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("values: Tvalues") 173dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tkeys: type") 174dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tvalues: type") 175dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 176dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 177dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle)); 178dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 179dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle values = c->UnknownShape(); 180dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRankAtLeast(values, 1, &values)); 181dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle keys = c->Vector(c->Dim(values, 0)); 182dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon c->set_output(0, keys); 183dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon c->set_output(1, values); 184dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 1853e852d462aaba446f62f76007405c0794a6087b9Anna R }); 186dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 187dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("LookupTableImport") 188dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: Ref(string)") 189dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("keys: Tin") 190dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("values: Tout") 191dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tin: type") 192dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tout: type") 193dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 194dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 195dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &handle)); 196dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon DimensionHandle unused_dim; 197dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithValue(c->Dim(handle, 0), 2, &unused_dim)); 198dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 199dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon // TODO(ebrevdo): Validate keys and values shape. 200dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 2013e852d462aaba446f62f76007405c0794a6087b9Anna R }); 202dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 203dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("LookupTableImportV2") 204dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: resource") 205dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("keys: Tin") 206dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("values: Tout") 207dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tin: type") 208dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tout: type") 209dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 210dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 211dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle)); 212dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 213dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon // TODO: Validate keys and values shape. 214dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 2153e852d462aaba446f62f76007405c0794a6087b9Anna R }); 216dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 217dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("HashTable") 218dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("table_handle: Ref(string)") 219dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("container: string = ''") 220dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("shared_name: string = ''") 221dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("use_node_name_sharing: bool = false") 222dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("key_dtype: type") 223dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_dtype: type") 224dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetIsStateful() 2253e852d462aaba446f62f76007405c0794a6087b9Anna R .SetShapeFn(TwoElementOutput); 226dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 227dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("HashTableV2") 228dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("table_handle: resource") 229dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("container: string = ''") 230dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("shared_name: string = ''") 231dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("use_node_name_sharing: bool = false") 232dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("key_dtype: type") 233dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_dtype: type") 234dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetIsStateful() 2353e852d462aaba446f62f76007405c0794a6087b9Anna R .SetShapeFn(ScalarOutput); 236dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 237dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("MutableHashTable") 238dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("table_handle: Ref(string)") 239dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("container: string = ''") 240dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("shared_name: string = ''") 241dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("use_node_name_sharing: bool = false") 242dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("key_dtype: type") 243dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_dtype: type") 244dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetIsStateful() 2453e852d462aaba446f62f76007405c0794a6087b9Anna R .SetShapeFn(TwoElementOutput); 246dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 247dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("MutableHashTableV2") 248dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("table_handle: resource") 249dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("container: string = ''") 250dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("shared_name: string = ''") 251dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("use_node_name_sharing: bool = false") 252dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("key_dtype: type") 253dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_dtype: type") 254dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetIsStateful() 2553e852d462aaba446f62f76007405c0794a6087b9Anna R .SetShapeFn(ScalarOutput); 256dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 257dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("MutableHashTableOfTensors") 258dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("table_handle: Ref(string)") 259dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("container: string = ''") 260dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("shared_name: string = ''") 261dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("use_node_name_sharing: bool = false") 262dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("key_dtype: type") 263dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_dtype: type") 264dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_shape: shape = {}") 265dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetIsStateful() 2663e852d462aaba446f62f76007405c0794a6087b9Anna R .SetShapeFn(TwoElementOutput); 267dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 268dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("MutableHashTableOfTensorsV2") 269dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("table_handle: resource") 270dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("container: string = ''") 271dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("shared_name: string = ''") 272dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("use_node_name_sharing: bool = false") 273dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("key_dtype: type") 274dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_dtype: type") 275dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_shape: shape = {}") 276dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetIsStateful() 2773e852d462aaba446f62f76007405c0794a6087b9Anna R .SetShapeFn(ScalarOutput); 278dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 279dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("MutableDenseHashTable") 280dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("empty_key: key_dtype") 281dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("table_handle: Ref(string)") 282dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("container: string = ''") 283dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("shared_name: string = ''") 284dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("use_node_name_sharing: bool = false") 285dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("key_dtype: type") 286dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_dtype: type") 287dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_shape: shape = {}") 288dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("initial_num_buckets: int = 131072") // 2^17 289dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("max_load_factor: float = 0.8") 290dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetIsStateful() 2913e852d462aaba446f62f76007405c0794a6087b9Anna R .SetShapeFn(TwoElementOutput); 292dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 293dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("MutableDenseHashTableV2") 294dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("empty_key: key_dtype") 295dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Output("table_handle: resource") 296dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("container: string = ''") 297dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("shared_name: string = ''") 298dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("use_node_name_sharing: bool = false") 299dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("key_dtype: type") 300dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_dtype: type") 301dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_shape: shape = {}") 302dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("initial_num_buckets: int = 131072") // 2^17 303dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("max_load_factor: float = 0.8") 304dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetIsStateful() 3053e852d462aaba446f62f76007405c0794a6087b9Anna R .SetShapeFn(ScalarOutput); 306dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 307dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("InitializeTable") 308dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: Ref(string)") 309dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("keys: Tkey") 310dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("values: Tval") 311dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tkey: type") 312dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tval: type") 313dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 314dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 315dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &handle)); 316dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon DimensionHandle unused_dim; 317dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithValue(c->Dim(handle, 0), 2, &unused_dim)); 318dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 319dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle keys; 320dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &keys)); 321dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->Merge(keys, c->input(2), &keys)); 322dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 3233e852d462aaba446f62f76007405c0794a6087b9Anna R }); 324dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 325dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("InitializeTableV2") 326dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: resource") 327dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("keys: Tkey") 328dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("values: Tval") 329dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tkey: type") 330dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("Tval: type") 331dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 332dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 333dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle)); 334dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 335dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle keys; 336dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &keys)); 337dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->Merge(keys, c->input(2), &keys)); 338dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 3393e852d462aaba446f62f76007405c0794a6087b9Anna R }); 340dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 341dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("InitializeTableFromTextFile") 342dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: Ref(string)") 343dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("filename: string") 344dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("key_index: int >= -2") 345dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_index: int >= -2") 346dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("vocab_size: int >= -1 = -1") 347dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("delimiter: string = '\t'") 348dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 349dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 350dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &handle)); 351dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon DimensionHandle unused_dim; 352dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithValue(c->Dim(handle, 0), 2, &unused_dim)); 353dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 354dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &handle)); 355dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 3563e852d462aaba446f62f76007405c0794a6087b9Anna R }); 357dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 358dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka LeonREGISTER_OP("InitializeTableFromTextFileV2") 359dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("table_handle: resource") 360dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Input("filename: string") 361dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("key_index: int >= -2") 362dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("value_index: int >= -2") 363dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("vocab_size: int >= -1 = -1") 364dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .Attr("delimiter: string = '\t'") 365dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon .SetShapeFn([](InferenceContext* c) { 366dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon ShapeHandle handle; 367dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle)); 368dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 369dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &handle)); 370dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon return Status::OK(); 3713e852d462aaba446f62f76007405c0794a6087b9Anna R }); 372dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon 373dd140f79e06a81c52cd8fc9ec6cda975a78a401fYutaka Leon} // namespace tensorflow 374