}

  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);

  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 =

  if (!func) {
    return emitError("'f' attribute refers to an undefined function: ")
           << func_name;
  }

  FunctionType func_ty = func.getFunctionType();
  int func_arg_count = func_ty.getNumInputs();
  int arg_count = getArgs().size();

  if (arg_count != func_arg_count) {
    return emitError() << "argument count mismatch: 'args' has " << arg_count