Value max_output_size = func_.getArgument(2);
  Value iou_threshold = func_.getArgument(3);
  Value score_threshold = func_.getArgument(4);
  auto output_type0 = func_.getFunctionType().getResult(0);
  auto output_type1 = func_.getFunctionType().getResult(1);

  OpBuilder builder(func_.getBody());
  auto op = builder.create<mlir::TFL::NonMaxSuppressionV4Op>(
      func_.getLoc(), output_type0, output_type1, boxes, scores,

           << func_.getNumArguments();
  }
  if (func_.getFunctionType().getNumResults() != 1) {
    return func_.emitWarning()
           << "Invalid number of results from " << kImageWarping << ": "
           << func_.getFunctionType().getNumResults();
  }

  // Check types and shapes.
  auto image_type = mlir::dyn_cast_or_null<RankedTensorType>(
      func_.getFunctionType().getInput(0));

  // * Zero or more row_split tensors.
  constexpr int kValues = 0;
  constexpr int kRowSplits = 1;

  if (func.getFunctionType().getInputs().size() !=
      func.getFunctionType().getResults().size()) {
    return func.emitError() << "Mismatched number of inputs and outputs.";
  }

  int row_splits = func.getFunctionType().getInputs().size() - kRowSplits;
  if (row_splits == 0) {
    auto input_values = GetInputType(func, kValues);

                                   .getValue())
          .getValues<FloatAttr>()[0]
          .getValue()
          .isExactlyValue(0.0f));

  EXPECT_EQ(fused_lstm_func_.getFunctionType().getNumResults(), 1);
  auto output_types = fused_lstm_func_.getFunctionType().getResults();
  SmallVector<int64_t, 2> output_shape{1, mlir::ShapedType::kDynamic};
  EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getShape().size(),

  if (op->getNumOperands() != func.getFunctionType().getNumInputs())
    return op.emitError()
           << "'host_func' has " << func.getFunctionType().getNumInputs()
           << " inputs and '_XlaHostComputeMlir' has " << op->getNumOperands()
           << " operands.  Number of operands/inputs should be the same.";

  if (op->getNumResults() != func.getFunctionType().getNumResults())

void init_builders(py::module& m) {
  py::class_<mlir::Builder>(m, "Builder")
      .def(py::init<mlir::MLIRContext*>())
      .def("getFunctionType",
           [](mlir::Builder& b, std::vector<mlir::Type> inputs,
              std::vector<mlir::Type> outputs) {
             return b.getFunctionType(llvm::ArrayRef<mlir::Type>(inputs),
                                      llvm::ArrayRef<mlir::Type>(outputs));
           });

      return;
    }
    if (main.getFunctionType().getNumResults() != op->getNumResults()) {
      op->emitError() << "mismatch between MLIR Opaque Op number of results ("
                      << op->getNumResults()
                      << ") and main() entry point in the module ("
                      << main.getFunctionType().getNumResults()
                      << " results)\n";
      return;
    }

  OpBuilder builder(&body, body.begin());
  auto graph_op = builder.create<tf_executor::GraphOp>(
      loc, func.getFunctionType().getResults());
  graph_op.getBody().push_back(new Block);
  builder.setInsertionPointToEnd(&graph_op.GetBody());
  auto island = builder.create<tf_executor::IslandOp>(
      loc, func.getFunctionType().getResults(),
      tf_executor::ControlType::get(&getContext()), ArrayRef<Value>());
  // Create Fetch.

  llvm::SmallVector<Type, 4> operand_types;
  operand_types.reserve(live_ins.size());
  for (Value v : live_ins) operand_types.emplace_back(v.getType());

  auto func_type = builder->getFunctionType(operand_types, op.getResultTypes());

  std::string func_name;
  if (auto outlined_func_name = op->template getAttrOfType<StringAttr>(
          TF::kClusterOutlinedFunctionNameAttr)) {

  auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
  if (!main_fn) {
    return absl::InternalError("Could not find main function in MLIR Module.");
  }

  mlir::FunctionType func_type = main_fn.getFunctionType();
  for (auto input_type : func_type.getInputs()) {
    tensorflow::TensorShape tensor_shape;
    xla::Shape xla_shape = xla::TypeToShape(input_type);
    TF_RETURN_IF_ERROR(tensorflow::TensorShape::BuildTensorShape(