1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_SHAPE_INFERENCE_UTILS_H_ 17 #define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_SHAPE_INFERENCE_UTILS_H_ 18 19 #include "llvm/ADT/Optional.h" 20 #include "tensorflow/core/ir/utils/shape_inference_utils.h" 21 22 namespace mlir { 23 24 class Operation; 25 26 namespace TF { 27 28 // Runs TensorFlow shape inference associated to the op type registered in the 29 // TensorFlow op registry based on the Graph version, operands, and attributes. 30 // Invoking this shape function will create conversions of parameters to the 31 // TensorFlow Graph equivalent data structures and back to MLIR equivalent data 32 // structures. This does not use a natively implemented shape inference in MLIR, 33 // and instead is temporary until shape functions are reimplemented/migrated to 34 // being in MLIR instead of the TensorFlow op registry. 35 LogicalResult InferReturnTypeComponentsForTFOp( 36 Optional<Location> location, Operation* op, int64_t graph_version, 37 tfg::OperandAsConstantFn operand_as_constant_fn, 38 tfg::OpResultAsShapeFn op_result_as_shape_fn, 39 tfg::ResultElementTypeFn result_element_type_fn, 40 SmallVectorImpl<ShapedTypeComponents>& inferred_return_shapes); 41 42 } // namespace TF 43 } // namespace mlir 44 45 #endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_SHAPE_INFERENCE_UTILS_H_ 46