op_lib_names = [
        "array_ops",
    ],
    pkg = "//tensorflow/core",
)

tf_gen_op_wrappers_cc(
    name = "dataset_ops",
    api_def_srcs = ["//tensorflow/core/api_def:base_api_def"],
    op_lib_names = [
        "dataset_ops",
        "experimental_dataset_ops",
    ],
    pkg = "//tensorflow/core",
)

tf_gen_op_wrappers_cc(
    name = "cc_ops",

from tensorflow.core.framework import tensor_shape_pb2
from tensorflow.core.framework import types_pb2
from tensorflow.python.client import session
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import def_function
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape

  ArrayAttr shapes;
  ArrayAttr types;
};

// Returns the input elements shapes and types for Dataset ops.
DatasetInput GetDatasetInput(Value value) {
  // TODO(haoliang): add an interface for DatasetOp to avoid the following
  // enumeration.
  // Iteratively tracing upwards if parent op is `IdentityOp` or `IdentityNOp`.
  while (
      llvm::isa_and_nonnull<IdentityOp, IdentityNOp>(value.getDefiningOp())) {