"contain 'host_func' function.";

  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.";

  // TODO(jpienaar): Expand to support inputs.
  mlir::func::FuncOp main = module->lookupSymbol<mlir::func::FuncOp>("main");
  QCHECK(main) << "No 'main' function specified.";
  if (main.getFunctionType().getNumInputs() != 0)
    LOG(QFATAL) << "NYI: Only nullary functions supported.";

  // Convert to flatbuffer.
  std::string serialized_flatbuffer;
  tflite::FlatbufferExportOptions options;

  virtual llvm::StringRef GetCompositeOpName() { return kLstmCellSimple; }

  // Rewrite the func body with constructed fused lstm.
  LogicalResult RewriteFunc();

  int GetNumInputs() { return n_input_; }

 protected:
  // verify input func op arguments/attributes and initialize internal state.
  virtual LogicalResult InitializeFromFuncAttributes();
  virtual LogicalResult Initialize();

  }

  auto func_type = func.getFunctionType();

  if (func_type.getNumInputs() != 2) {
    return op.emitOpError()
           << "expects reduction function to take 2 parameters, but "
              "has "
           << func_type.getNumInputs() << " parameter(s)";
  }

  return success();
}

  // input and the outputs are all using this type and predicate is tensor<i1>
  // type).
  if (mlir::isa<FunctionType>(types.front())) {
    FunctionType type = mlir::cast<FunctionType>(types.front());
    if (type.getNumInputs() < 2)
      return parser.emitError(parser.getNameLoc())
             << " expects a single data type and a predicate";
    result.types.assign(type.getResults().begin(), type.getResults().end());

  builder.create<func::ReturnOp>(func.getLoc(), op.getResults());
  return success();
}

LogicalResult VerifySgnnProjection(func::FuncOp func, FuncAttr attr) {
  if (func.getFunctionType().getNumInputs() != 2 ||
      func.getFunctionType().getNumResults() != 1) {
    return func.emitError() << "Mismatched number of inputs and outputs.";
  }
  auto values_type = GetInputType(func, 0);

                         func::FuncOp fn, OpBuilder* builder) {
  FunctionType fn_type = fn.getFunctionType();
  llvm::SmallVector<Value, 4> operands;
  int num_operands = fn_type.getNumInputs();
  operands.reserve(num_operands);
  for (int i = 0; i < num_operands; ++i) {
    Value val = get_arg(i);
    Type expected = fn_type.getInput(i);
    if (val.getType() != expected) {
      val =

  for (func::FuncOp func : {op.cond_function(), op.body_function()}) {
    ++func_index;
    if (!func) continue;

    FunctionType func_type = func.getFunctionType();
    int num_inputs = func_type.getNumInputs();
    int num_results = func_type.getNumResults();

    // For each argument type in function's arguments, change it to uranked
    // tensor type if it's a variant type.

      module.lookupSymbol<mlir::func::FuncOp>("main");
  TF_RET_CHECK(main_func != nullptr) << "No main function found";
  mlir::FunctionType func_type = main_func.getFunctionType();

  int num_args = func_type.getNumInputs();
  xla_input_shapes->reserve(num_args);

  std::vector<xla::Shape> individual_arg_shapes;
  individual_arg_shapes.reserve(num_args);
  for (int i = 0; i < num_args; ++i) {

  FunctionType func_type = func.getFunctionType();
  bool needs_refinement = false;
  SmallVector<Type, 4> new_arg_types;
  new_arg_types.reserve(func_type.getNumInputs());

  // Update argument types in-place using the provided arg_shapes.
  for (size_t i = 0; i < func_type.getNumInputs(); ++i) {
    ArrayRef<int64_t> shape = arg_shapes[i];
    Type element_type;
    if (auto input_ty =