};

ShapeInference::ShapeInference(int64_t graph_version, ModuleOp module,
                               bool propagate_caller_callee_constants,
                               ArrayRef<TypeID> ops_to_skip)
    : tf_dialect_(module->getContext()->getLoadedDialect<TensorFlowDialect>()),
      symbol_users_(symbol_table_, module),
      graph_version_(graph_version),

#define DEBUG_TYPE "tf-shape-inference-utils"

namespace mlir {
namespace TF {

LogicalResult InferReturnTypeComponentsForTFOp(
    std::optional<Location> location, Operation* op, int64_t graph_version,
    tfg::OperandAsConstantFn operand_as_constant_fn,
    tfg::OpResultAsShapeFn op_result_as_shape_fn,
    tfg::ResultElementTypeFn result_element_type_fn,
    SmallVectorImpl<ShapedTypeComponents>& inferred_return_shapes) {

// being in MLIR instead of the TensorFlow op registry.
LogicalResult InferReturnTypeComponentsForTFOp(
    std::optional<Location> location, Operation* op, int64_t graph_version,
    tfg::OperandAsConstantFn operand_as_constant_fn,
    tfg::OpResultAsShapeFn op_result_as_shape_fn,
    tfg::ResultElementTypeFn result_element_type_fn,
    SmallVectorImpl<ShapedTypeComponents>& inferred_return_shapes);

// reached convergence, false otherwise.
FailureOr<bool> InferShapeForFunction(func::FuncOp func,
                                      ArrayRef<ArrayRef<int64_t>> arg_shapes,
                                      int64_t graph_version,
                                      int64_t max_iterations = 10,
                                      ArrayRef<TypeID> ops_to_skip = {});

}  // namespace TF

}  // namespace mlir