| 1058 | } // namespace tensorflow |
| 1059 | |
| 1060 | void TFE_InferShapes(TFE_Op* tfe_op, TF_ShapeAndTypeList* input_shapes, |
| 1061 | TF_Tensor** input_tensors, |
| 1062 | TF_ShapeAndTypeList* input_tensors_as_shapes, |
| 1063 | TF_ShapeAndTypeList** input_resource_shapes_and_types, |
| 1064 | TF_ShapeAndTypeList** output_shapes, |
| 1065 | TF_ShapeAndTypeList*** output_resource_shapes_and_types, |
| 1066 | TF_Status* status) { |
| 1067 | using tensorflow::NodeDef; |
| 1068 | using tensorflow::OpRegistrationData; |
| 1069 | using tensorflow::Tensor; |
| 1070 | using tensorflow::shape_inference::DimensionHandle; |
| 1071 | using tensorflow::shape_inference::InferenceContext; |
| 1072 | using tensorflow::shape_inference::ShapeAndType; |
| 1073 | using tensorflow::shape_inference::ShapeHandle; |
| 1074 | |
| 1075 | const int num_inputs = input_shapes->num_items; |
| 1076 | NodeDef node_def; |
| 1077 | node_def.set_name(tfe_op->operation.Name()); |
| 1078 | node_def.set_op(tfe_op->operation.Name()); |
| 1079 | for (int i = 0; i < num_inputs; ++i) { |
| 1080 | node_def.add_input("dummy_input"); |
| 1081 | } |
| 1082 | tfe_op->operation.Attrs().FillAttrValueMap(node_def.mutable_attr()); |
| 1083 | |
| 1084 | const tensorflow::OpRegistrationData* op_reg_data; |
| 1085 | status->status = |
| 1086 | tensorflow::OpRegistry::Global()->LookUp(node_def.op(), &op_reg_data); |
| 1087 | if (!status->status.ok()) return; |
| 1088 | |
| 1089 | // Initialize a input_tensor vector with `nullptr` values. |
| 1090 | std::vector<const Tensor*> input_tensors_vector(num_inputs, nullptr); |
| 1091 | // A vector to keep track of newly created `tf::Tensor` objects. |
| 1092 | std::vector<Tensor> all_input_tensors; |
| 1093 | // Update the vector with information from `input_tensors` if provided. |
| 1094 | if (input_tensors != nullptr) { |
| 1095 | // Note that we take the address of the elements in `all_input_tensors` |
| 1096 | // below. Allocate enough space so that no reallocation happens, which will |
| 1097 | // make the pointers invalid. |
| 1098 | all_input_tensors.reserve(num_inputs); |
| 1099 | for (int i = 0; i < num_inputs; ++i) { |
| 1100 | if (input_tensors[i] == nullptr) continue; |
| 1101 | all_input_tensors.emplace_back(); |
| 1102 | Tensor& input_tensor = all_input_tensors.back(); |
| 1103 | status->status = TF_TensorToTensor(input_tensors[i], &input_tensor); |
| 1104 | if (!status->status.ok()) return; |
| 1105 | input_tensors_vector[i] = &input_tensor; |
| 1106 | } |
| 1107 | } |
| 1108 | |
| 1109 | // Create an inference context with dummy values, which will be updated later. |
| 1110 | InferenceContext c(TF_GRAPH_DEF_VERSION, &node_def, op_reg_data->op_def, |
| 1111 | std::vector<ShapeHandle>(num_inputs), input_tensors_vector, |
| 1112 | {}, |
| 1113 | std::vector<std::unique_ptr<std::vector<ShapeAndType>>>()); |
| 1114 | |
| 1115 | // Set input_shapes. |
| 1116 | for (int i = 0; i < num_inputs; ++i) { |
| 1117 | std::vector<DimensionHandle> dims; |