From 736e055a756cf0f99d59b67284aade01baec9799 Mon Sep 17 00:00:00 2001 From: Skye Wanderman-Milne Date: Tue, 27 Mar 2018 15:07:05 -0700 Subject: Make _USE_C_API = True and _USE_C_SHAPES = False work with handle data, take 2. This change makes _set_shapes_for_outputs_c_api fetch and set Tensor._handle_data. This is necessary for running the Python shape inference code on resource tensors. PiperOrigin-RevId: 190681459 --- tensorflow/c/python_api.cc | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) (limited to 'tensorflow/c/python_api.cc') diff --git a/tensorflow/c/python_api.cc b/tensorflow/c/python_api.cc index cd604538f1..93155998b8 100644 --- a/tensorflow/c/python_api.cc +++ b/tensorflow/c/python_api.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/c/python_api.h" #include "tensorflow/c/c_api_internal.h" +#include "tensorflow/python/framework/cpp_shape_inference.pb.h" namespace tensorflow { @@ -109,4 +110,29 @@ void ExtendSession(TF_Session* session, TF_Status* status) { session->extend_before_run = false; } +std::string ResourceHandleShapeAndType(TF_Graph* graph, TF_Output output) { + Node* node = &output.oper->node; + CppShapeInferenceResult::HandleData handle_data; + handle_data.set_is_set(true); + { + mutex_lock l(graph->mu); + tensorflow::shape_inference::InferenceContext* ic = + graph->refiner.GetContext(node); + CHECK(ic != nullptr); + CHECK_LT(output.index, ic->num_outputs()); + const auto* shapes_and_types = + ic->output_handle_shapes_and_types(output.index); + if (shapes_and_types == nullptr) return ""; + + for (const auto& p : *shapes_and_types) { + auto* out_shape_and_type = handle_data.add_shape_and_type(); + ic->ShapeHandleToProto(p.shape, out_shape_and_type->mutable_shape()); + out_shape_and_type->set_dtype(p.dtype); + } + } + string result; + handle_data.SerializeToString(&result); + return result; +} + } // namespace tensorflow -- cgit v1.2.3