LogicalResult matchAndRewrite(BatchDatasetV2Op op,
                                PatternRewriter &rewriter) const override {
    auto batchInputDataset = op.getInputDataset();

    ParallelMapDatasetOp batchInputOp = dyn_cast_or_null<ParallelMapDatasetOp>(
        batchInputDataset.getDefiningOp());
    if (!batchInputOp) return failure();

    // The type of the `num_parallel_calls` argument in ParallelMapDataset

        mlir::cast<OpResult>(value).getResultNumber());
  }

  Operation* op = value.getDefiningOp();
  if (!llvm::isa_and_nonnull<BatchDatasetV2Op, MapDatasetOp, RepeatDatasetOp,
                             ParallelMapDatasetOp, ParallelMapDatasetV2Op,
                             TakeDatasetOp, TakeWhileDatasetOp>(op))
    return DatasetInput{nullptr, nullptr};

  return DatasetInput{op->getAttrOfType<ArrayAttr>("output_shapes"),