if (op.getCond().front().getNumArguments() !=
      op.getBody().front().getNumArguments())
    return op.emitOpError(llvm::formatv(
        "number of arguments in condition function does not match number of "
        "arguments in body function ({0} != {1})",
        op.getCond().front().getNumArguments(),
        op.getBody().front().getNumArguments()));
  // Verify shapes are compatible.

  // 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())) {

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