}

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