// Store the arg/return attributes as a list rather than uniqueuing during
  // construction.
  llvm::SmallVector<mlir::NamedAttrList, 4> arg_attrs;
  arg_attrs.resize(func.getNumArguments());
  llvm::SmallVector<mlir::NamedAttrList, 4> ret_attrs;
  ret_attrs.resize(func.getNumResults());

  auto set_attributes_on_func = [&](Node* node, int64_t index, bool is_arg) {

}

bool ShapeInference::InferShapeForDatasetOpCommon(Operation* op, FuncOp f,
                                                  int64_t max_iterations) {
  int N = op->getNumOperands() - 1;
  int M = f.getNumArguments() - N;
  DCOMMENT_OP(op, "Inferring shape for with N = " << N << " and M = " << M);

  // Initialize with function input types.
  auto input_types = llvm::to_vector<1>(

          mlir::dyn_cast_or_null<mlir::StringAttr>(dict_attr.get("inputs"))) {
    str.getValue().split(input_names, ',', /*MaxSplit=*/-1,
                         /*KeepEmpty=*/false);
    if (input_names.size() != fn.getNumArguments()) {
      fn.emitWarning() << "invalid entry function specification";
      return;
    }
    for (const auto& it : llvm::enumerate(fn.getArguments())) {

      if (it != infos.end() && !it->getSecond().used) {
        return_op->eraseOperand(entry.index() - skipped_retvals++);
      }
    }
  }
  llvm::BitVector indices_to_erase(func_op.getNumArguments());
  llvm::SmallVector<Type, 4> new_types;
  int64_t skipped_args = 0;
  for (auto arg : func_op.getArguments()) {
    auto it = infos.find(arg.getArgNumber());
    if (it != infos.end() && !it->getSecond().used) {

                                   updated_result_types));
  });
  Region &entry = func.getRegion();
  TypeConverter::SignatureConversion signature_conversion(
      entry.getNumArguments());
  for (const BlockArgument &arg : entry.getArguments()) {
    signature_conversion.addInputs(arg.getArgNumber(),
                                   updated_argument_types[arg.getArgNumber()]);
  }

 public:
  using OpRewritePattern<stablehlo::ReduceWindowOp>::OpRewritePattern;

  LogicalResult MatchBinaryReduceFunction(Region& function) const {
    Block& body = function.front();
    if (body.getNumArguments() != 2) return failure();

    auto return_op = dyn_cast<stablehlo::ReturnOp>(body.back());
    if (!return_op) return failure();
    if (return_op.getNumOperands() != 1) return failure();