}

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

  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