/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/full_type.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def_builder.h"
#include "tensorflow/core/framework/shape_inference.h"

namespace tensorflow {

// --------------------------------------------------------------------------

// The ops in this section can be composed to define an input
// pipeline. Each op produces a DT_VARIANT tensor that represents
// a DAG of "dataset" objects. An "dataset" object can be converted
// to a stateful "iterator" by passing the "dataset" to the
// "MakeIterator" op.
//
// TODO(b/123753214): DT_VARIANT tensors that represent "dataset" objects are
// not presently serializable. To avoid issues with graph optimizations, such
// as constant folding, CSE, or DCE, ensure that any "source dataset" ops
// (i.e. ops that output a dataset and do not take one as input) are
// marked as "do not optimize".

// TODO(mrry): Validate that `components` have shapes compatible with
// `output_shapes`.
REGISTER_OP("TensorDataset")
    .Input("components: Toutput_types")
    .Output("handle: variant")
    .Attr("Toutput_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetDoNotOptimize()  // TODO(b/123753214): See comment in dataset_ops.cc.
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "Toutput_types"))
    .SetShapeFn(shape_inference::ScalarShape);

// TODO(mrry): Validate that the dim-0 slices of `components` have shapes
// compatible with `output_shapes`.
REGISTER_OP("TensorSliceDataset")
    .Input("components: Toutput_types")
    .Output("handle: variant")
    .Attr("Toutput_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("is_files: bool = false")
    .Attr("metadata: string = ''")
    .Attr("replicate_on_split: bool = false")
    .SetDoNotOptimize()  // TODO(b/123753214): See comment in dataset_ops.cc.
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "Toutput_types"))
    .SetForwardTypeFn(full_type::MultiaryUnstack(TFT_DATASET,
                                                 full_type::UnstackTensor))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("SparseTensorSliceDataset")
    .Input("indices: int64")
    .Input("values: Tvalues")
    .Input("dense_shape: int64")
    .Output("handle: variant")
    .Attr("Tvalues: type")
    .SetDoNotOptimize()  // TODO(b/123753214): See comment in dataset_ops.cc.
    .SetTypeConstructor(full_type::UnaryTensorContainer(TFT_DATASET, "Tvalues"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("GeneratorDataset")
    .Input("init_func_other_args: Tinit_func_args")
    .Input("next_func_other_args: Tnext_func_args")
    .Input("finalize_func_other_args: Tfinalize_func_args")
    .Output("handle: variant")
    .Attr("init_func: func")
    .Attr("next_func: func")
    .Attr("finalize_func: func")
    .Attr("Tinit_func_args: list(type) >= 0")
    .Attr("Tnext_func_args: list(type) >= 0")
    .Attr("Tfinalize_func_args: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetDoNotOptimize()  // TODO(b/123753214): See comment in dataset_ops.cc.
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("ZipDataset")
    .Input("input_datasets: N * variant")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("N: int >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("ConcatenateDataset")
    .Input("input_dataset: variant")
    .Input("another_dataset: variant")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("RepeatDataset")
    .Input("input_dataset: variant")
    .Input("count: int64")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle count_shape;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &count_shape));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("TakeDataset")
    .Input("input_dataset: variant")
    .Input("count: int64")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle count_shape;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &count_shape));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("SkipDataset")
    .Input("input_dataset: variant")
    .Input("count: int64")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle count_shape;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &count_shape));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("MapDataset")
    .Input("input_dataset: variant")
    .Input("other_arguments: Targuments")
    .Output("handle: variant")
    .Attr("f: func")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("use_inter_op_parallelism: bool = true")
    .Attr("preserve_cardinality: bool = false")
    .Attr("force_synchronous: bool = false")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("ParallelMapDataset")
    .Input("input_dataset: variant")
    .Input("other_arguments: Targuments")
    .Input("num_parallel_calls: int32")
    .Output("handle: variant")
    .Attr("f: func")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("use_inter_op_parallelism: bool = true")
    .Attr("sloppy: bool = false")
    .Attr("preserve_cardinality: bool = false")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("ParallelMapDatasetV2")
    .Input("input_dataset: variant")
    .Input("other_arguments: Targuments")
    .Input("num_parallel_calls: int64")
    .Output("handle: variant")
    .Attr("f: func")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("use_inter_op_parallelism: bool = true")
    // "true", "false", or "default".
    .Attr("deterministic: string = 'default'")
    .Attr("preserve_cardinality: bool = false")
    .Attr("use_unbounded_threadpool: bool = false")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("PrefetchDataset")
    .Input("input_dataset: variant")
    .Input("buffer_size: int64")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("slack_period: int = 0")
    .Attr("legacy_autotune: bool = true")
    .Attr("buffer_size_min: int = 0")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // buffer_size should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("FlatMapDataset")
    .Input("input_dataset: variant")
    .Input("other_arguments: Targuments")
    .Output("handle: variant")
    .Attr("f: func")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("InterleaveDataset")
    .Input("input_dataset: variant")
    .Input("other_arguments: Targuments")
    .Input("cycle_length: int64")
    .Input("block_length: int64")
    .Output("handle: variant")
    .Attr("f: func")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("ParallelInterleaveDatasetV2")
    .Input("input_dataset: variant")
    .Input("other_arguments: Targuments")
    .Input("cycle_length: int64")
    .Input("block_length: int64")
    .Input("num_parallel_calls: int64")
    .Output("handle: variant")
    .Attr("f: func")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("sloppy: bool = false")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("ParallelInterleaveDatasetV3")
    .Input("input_dataset: variant")
    .Input("other_arguments: Targuments")
    .Input("cycle_length: int64")
    .Input("block_length: int64")
    .Input("num_parallel_calls: int64")
    .Output("handle: variant")
    .Attr("f: func")
    // "true", "false", or "default".
    .Attr("deterministic: string = 'default'")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

// Like V3, but adds buffer_output_elements and prefetch_input_elements.
REGISTER_OP("ParallelInterleaveDatasetV4")
    .Input("input_dataset: variant")
    .Input("other_arguments: Targuments")
    .Input("cycle_length: int64")
    .Input("block_length: int64")
    .Input("buffer_output_elements: int64")
    .Input("prefetch_input_elements: int64")
    .Input("num_parallel_calls: int64")
    .Output("handle: variant")
    .Attr("f: func")
    // "true", "false", or "default".
    .Attr("deterministic: string = 'default'")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("FilterDataset")
    .Input("input_dataset: variant")
    .Input("other_arguments: Targuments")
    .Output("handle: variant")
    .Attr("predicate: func")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("DatasetFingerprint")
    .Input("input_dataset: variant")
    .Output("fingerprint: uint64")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("ParallelFilterDataset")
    .Input("input_dataset: variant")
    .Input("other_arguments: Targuments")
    .Input("num_parallel_calls: int64")
    .Output("handle: variant")
    .Attr("predicate: func")
    // "true", "false", or "default".
    .Attr("deterministic: string = 'default'")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

// This op is no longer supported.
REGISTER_OP("FilterByLastComponentDataset")
    .Input("input_dataset: variant")
    .Output("output: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("WindowDataset")
    .Input("input_dataset: variant")
    .Input("size: int64")
    .Input("shift: int64")
    .Input("stride: int64")
    .Input("drop_remainder: bool")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // size, shift, stride, and drop_remainder should be scalars.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("WindowOp")
    .Input("inputs: Tinputs")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("Tinputs: list(type) >= 1")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("BatchDataset")
    .Input("input_dataset: variant")
    .Input("batch_size: int64")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // batch_size should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("BatchDatasetV2")
    .Input("input_dataset: variant")
    .Input("batch_size: int64")
    .Input("drop_remainder: bool")
    .Output("handle: variant")
    .Attr("parallel_copy: bool = false")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetForwardTypeFn(full_type::ContainerMap(TFT_DATASET, /*input_idx=*/0,
                                              full_type::BatchTensor))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // batch_size should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      // drop_remainder should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("ParallelBatchDataset")
    .Input("input_dataset: variant")
    .Input("batch_size: int64")
    .Input("num_parallel_calls: int64")
    .Input("drop_remainder: bool")
    .Output("handle: variant")
    .Attr("parallel_copy: bool = false")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    // "true", "false", or "default".
    .Attr("deterministic: string = 'default'")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // batch_size should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      // num_parallel_calls should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      // drop_remainder should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("ShardDataset")
    .Input("input_dataset: variant")
    .Input("num_shards: int64")
    .Input("index: int64")
    .Output("handle: variant")
    .Attr("require_non_empty: bool = false")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // num_shards should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      // index should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

// TODO(mrry): Validate that `padded_shapes` are all vectors, the lengths of
// `output_types` and `output_shapes` are `N` the `output_shapes` are (as far as
// possible to tell statically) compatible with `padded_shapes`, and that
// `padding_values` are all scalars.
REGISTER_OP("PaddedBatchDataset")
    .Input("input_dataset: variant")
    .Input("batch_size: int64")
    .Input("padded_shapes: N * int64")
    .Input("padding_values: Toutput_types")
    .Output("handle: variant")
    .Attr("Toutput_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("N: int >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "Toutput_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // batch_size should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("PaddedBatchDatasetV2")
    .Input("input_dataset: variant")
    .Input("batch_size: int64")
    .Input("padded_shapes: N * int64")
    .Input("padding_values: Toutput_types")
    .Input("drop_remainder: bool")
    .Output("handle: variant")
    .Attr("parallel_copy: bool = false")
    .Attr("Toutput_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("N: int >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "Toutput_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // batch_size should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      // drop_remainder should be a scalar.
      TF_RETURN_IF_ERROR(
          c->WithRank(c->input(c->num_inputs() - 1), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("RangeDataset")
    .Input("start: int64")
    .Input("stop: int64")
    .Input("step: int64")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .Attr("replicate_on_split: bool = false")
    .SetDoNotOptimize()  // TODO(b/123753214): See comment in dataset_ops.cc.
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // start, stop, and step should be scalars.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("RewriteDataset")
    .Input("input_dataset: variant")
    .Input("rewrite_name: string")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("AnonymousSeedGenerator")
    .Input("seed: int64")
    .Input("seed2: int64")
    .Input("reshuffle: bool")
    .Output("handle: resource")
    .Output("deleter: variant")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      c->set_output(1, c->Scalar());
      return absl::OkStatus();
    });

REGISTER_OP("DatasetCardinality")
    .Input("input_dataset: variant")
    .Output("cardinality: int64")
    .Attr("cardinality_options: string = ''")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("DeleteSeedGenerator")
    .Input("handle: resource")
    .Input("deleter: variant")
    .SetShapeFn(shape_inference::NoOutputs);

// Deprecated in favor of AnonymousSeedGenerator/DeleteSeedGenerator.
REGISTER_OP("AnonymousRandomSeedGenerator")
    .Input("seed: int64")
    .Input("seed2: int64")
    .Output("handle: resource")
    .Output("deleter: variant")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      c->set_output(1, c->Scalar());
      return absl::OkStatus();
    });

// Deprecated in favor of AnonymousSeedGenerator/DeleteSeedGenerator.
REGISTER_OP("DeleteRandomSeedGenerator")
    .Input("handle: resource")
    .Input("deleter: variant")
    .SetShapeFn(shape_inference::NoOutputs);

REGISTER_OP("DummySeedGenerator")
    .Output("handle: resource")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      return absl::OkStatus();
    });

REGISTER_OP("ShuffleDataset")
    .Input("input_dataset: variant")
    .Input("buffer_size: int64")
    .Input("seed: int64")
    .Input("seed2: int64")
    .Output("handle: variant")
    .Attr("reshuffle_each_iteration: bool = true")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // buffer_size, seed, and seed2 should be scalars.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("ShuffleDatasetV2")
    .Input("input_dataset: variant")
    .Input("buffer_size: int64")
    .Input("seed_generator: resource")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // buffer_size and seed_generator should be scalars.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("ShuffleDatasetV3")
    .Input("input_dataset: variant")
    .Input("buffer_size: int64")
    .Input("seed: int64")
    .Input("seed2: int64")
    .Input("seed_generator: resource")
    .Output("handle: variant")
    .Attr("reshuffle_each_iteration: bool = true")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // buffer_size, seed, seed2, and seed_generator should be scalars.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("ShuffleAndRepeatDataset")
    .Input("input_dataset: variant")
    .Input("buffer_size: int64")
    .Input("seed: int64")
    .Input("seed2: int64")
    .Input("count: int64")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("reshuffle_each_iteration: bool = true")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // buffer_size, seed, seed2, and count should be scalars.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("ShuffleAndRepeatDatasetV2")
    .Input("input_dataset: variant")
    .Input("buffer_size: int64")
    .Input("seed: int64")
    .Input("seed2: int64")
    .Input("count: int64")
    .Input("seed_generator: resource")
    .Output("handle: variant")
    .Attr("reshuffle_each_iteration: bool = true")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // buffer_size, seed, seed2, count, and seed_generator should be scalars.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("AnonymousMemoryCache")
    .Output("handle: resource")
    .Output("deleter: variant")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      c->set_output(1, c->Scalar());
      return absl::OkStatus();
    });

REGISTER_OP("DeleteMemoryCache")
    .Input("handle: resource")
    .Input("deleter: variant")
    .SetShapeFn(shape_inference::NoOutputs);

REGISTER_OP("DummyMemoryCache")
    .Output("handle: resource")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      return absl::OkStatus();
    });

REGISTER_OP("CacheDataset")
    .Input("input_dataset: variant")
    .Input("filename: string")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    // TODO(mdan): Should these use type inference instead?
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // filename should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("CacheDatasetV2")
    .Input("input_dataset: variant")
    .Input("filename: string")
    .Input("cache: resource")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // filename should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      // cache should be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("TextLineDataset")
    .Input("filenames: string")
    .Input("compression_type: string")
    .Input("buffer_size: int64")
    .Attr("metadata: string = ''")
    .Output("handle: variant")
    .SetDoNotOptimize()  // TODO(b/123753214): See comment in dataset_ops.cc.
    .SetTypeConstructor(full_type::UnaryTensorContainer(TFT_DATASET,
                                                        TFT_STRING))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // `filenames` must be a scalar or a vector.
      TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused));
      // `compression_type` could only be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      // `buffer_size` could only be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("FixedLengthRecordDataset")
    .Input("filenames: string")
    .Input("header_bytes: int64")
    .Input("record_bytes: int64")
    .Input("footer_bytes: int64")
    .Input("buffer_size: int64")
    .Attr("metadata: string = ''")
    .Output("handle: variant")
    .SetDoNotOptimize()  // TODO(b/123753214): See comment in dataset_ops.cc.
    .SetTypeConstructor(full_type::UnaryTensorContainer(TFT_DATASET,
                                                        TFT_STRING))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // `filenames` must be a scalar or a vector.
      TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused));
      // header_bytes, record_bytes, footer_bytes, buffer_size should be
      // scalars.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("FixedLengthRecordDatasetV2")
    .Input("filenames: string")
    .Input("header_bytes: int64")
    .Input("record_bytes: int64")
    .Input("footer_bytes: int64")
    .Input("buffer_size: int64")
    .Input("compression_type: string")
    .Attr("metadata: string = ''")
    .Output("handle: variant")
    .SetDoNotOptimize()  // TODO(b/123753214): See comment in dataset_ops.cc.
    .SetTypeConstructor(full_type::UnaryTensorContainer(TFT_DATASET,
                                                        TFT_STRING))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // `filenames` must be a scalar or a vector.
      TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused));
      // header_bytes, record_bytes, footer_bytes, buffer_size should be
      // scalars.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused));
      TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("TFRecordDataset")
    .Input("filenames: string")
    .Input("compression_type: string")
    .Input("buffer_size: int64")
    .Attr("metadata: string = ''")
    .Output("handle: variant")
    .SetDoNotOptimize()  // TODO(b/123753214): See comment in dataset_ops.cc.
    .SetTypeConstructor(full_type::UnaryTensorContainer(TFT_DATASET,
                                                        TFT_STRING))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // `filenames` must be a scalar or a vector.
      TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused));
      // `compression_type` could only be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      // `buffer_size` could only be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("TFRecordDatasetV2")
    .Input("filenames: string")
    .Input("compression_type: string")
    .Input("buffer_size: int64")
    .Input("byte_offsets: int64")
    .Attr("metadata: string = ''")
    .Output("handle: variant")
    .SetDoNotOptimize()  // TODO(b/123753214): See comment in dataset_ops.cc.
    .SetTypeConstructor(full_type::UnaryTensorContainer(TFT_DATASET,
                                                        TFT_STRING))
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      shape_inference::ShapeHandle unused;
      // `filenames` must be a scalar or a vector.
      TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused));
      // `compression_type` could only be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
      // `buffer_size` could only be a scalar.
      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
      // `byte_offsets` must be a scalar or a vector.
      TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(3), 1, &unused));
      return shape_inference::ScalarShape(c);
    });

REGISTER_OP("Iterator")
    .Output("handle: resource")
    .Attr("shared_name: string")
    .Attr("container: string")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("IteratorV2")
    .Output("handle: resource")
    .Attr("shared_name: string")
    .Attr("container: string")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("AnonymousIterator")
    .Output("handle: resource")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("AnonymousIteratorV2")
    .Output("handle: resource")
    .Output("deleter: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      c->set_output(1, c->Scalar());
      return absl::OkStatus();
    });

REGISTER_OP("AnonymousIteratorV3")
    .Output("handle: resource")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      return absl::OkStatus();
    });

REGISTER_OP("DeleteIterator")
    .Input("handle: resource")
    .Input("deleter: variant")
    .SetShapeFn(shape_inference::NoOutputs);

REGISTER_OP("DeleteMultiDeviceIterator")
    .Input("multi_device_iterator: resource")
    .Input("iterators: N * resource")
    .Input("deleter: variant")
    .Attr("N: int >= 0")
    .SetShapeFn(shape_inference::NoOutputs);

REGISTER_OP("MakeIterator")
    .Input("dataset: variant")
    .Input("iterator: resource")
    .SetTypeConstructor(full_type::NoOutputs())
    .SetReverseTypeFn(1, full_type::MapCovariant(TFT_DATASET, TFT_ITERATOR, 0))
    .SetShapeFn(shape_inference::NoOutputs);

REGISTER_OP("OneShotIterator")
    .Output("handle: resource")
    .Attr("dataset_factory: func")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("container: string = ''")
    .Attr("shared_name: string = ''")
    .SetIsStateful()
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("IteratorGetNext")
    .Input("iterator: resource")
    .Output("components: output_types")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn(shape_inference::DatasetIteratorShape);

REGISTER_OP("IteratorGetNextSync")
    .Input("iterator: resource")
    .Output("components: output_types")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn(shape_inference::DatasetIteratorShape);

// TODO(b/124308596): Instead of conservatively marking this op as stateful,
// implement a mechanism to determine whether `dataset` has a side-effect
// and use it to decide whether to use a stateless or stateful version of this
// op.
REGISTER_OP("DatasetToSingleElement")
    .Input("dataset: variant")
    .Output("components: output_types")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetIsStateful()
    .SetShapeFn(shape_inference::DatasetIteratorShape);

// TODO(b/124308596): Instead of conservatively marking this op as stateful,
// implement a mechanism to determine whether `dataset` has a side-effect
// and use it to decide whether to use a stateless or stateful version of this
// op.
REGISTER_OP("ReduceDataset")
    .Input("input_dataset: variant")
    .Input("initial_state: Tstate")
    .Input("other_arguments: Targuments")
    .Output("components: output_types")
    .Attr("f: func")
    .Attr("Tstate: list(type) >= 1")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("use_inter_op_parallelism: bool = true")
    .Attr("metadata: string = ''")
    .SetIsStateful()
    .SetShapeFn(shape_inference::DatasetIteratorShape);

REGISTER_OP("IteratorToStringHandle")
    .Input("resource_handle: resource")
    .Output("string_handle: string")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("IteratorFromStringHandle")
    .Input("string_handle: string")
    .Output("resource_handle: resource")
    .Attr("output_types: list(type) >= 0 = []")
    .Attr("output_shapes: list(shape) >= 0 = []")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("IteratorFromStringHandleV2")
    .Input("string_handle: string")
    .Output("resource_handle: resource")
    .Attr("output_types: list(type) >= 0 = []")
    .Attr("output_shapes: list(shape) >= 0 = []")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("SerializeIterator")
    .Input("resource_handle: resource")
    .Attr("external_state_policy: int = 0")
    .Output("serialized: variant")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Vector(c->UnknownDim()));
      return absl::OkStatus();
    });

REGISTER_OP("DeserializeIterator")
    .Input("resource_handle: resource")
    .Input("serialized: variant")
    .SetShapeFn(shape_inference::NoOutputs);

REGISTER_OP("DatasetToGraph")
    .Input("input_dataset: variant")
    .Attr("stateful_whitelist: list(string) >= 0 = []")
    .Attr("allow_stateful: bool = false")
    .Attr("strip_device_assignment: bool = false")
    .Output("graph: string")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("DatasetToGraphV2")
    .Input("input_dataset: variant")
    .Attr("external_state_policy: int = 0")
    .Attr("strip_device_assignment: bool = false")
    .Output("graph: string")
    .SetForwardTypeFn(full_type::Encode(TFT_STRING, 0))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("OptimizeDataset")
    .Input("input_dataset: variant")
    .Input("optimizations: string")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("optimization_configs: list(string) = []")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("OptimizeDatasetV2")
    .Input("input_dataset: variant")
    .Input("optimizations_enabled: string")
    .Input("optimizations_disabled: string")
    .Input("optimizations_default: string")
    .Output("handle: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("optimization_configs: list(string) = []")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("IteratorGetNextAsOptional")
    .Input("iterator: resource")
    .Output("optional: variant")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_OPTIONAL,
                                                           "output_types"))
    .SetForwardTypeFn(full_type::MapCovariant(TFT_ITERATOR, TFT_OPTIONAL, 0))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("ModelDataset")
    .Input("input_dataset: variant")
    .Output("handle: variant")
    .Attr("algorithm: int = 0")
    .Attr("cpu_budget: int = 0")
    .Attr("ram_budget: int = 0")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

// TODO(b/124308749): Add a stateful version of MapDefun and use it when `f`
// is stateful.
REGISTER_OP("MapDefun")
    .Input("arguments: Targuments")
    .Input("captured_inputs: Tcaptured")
    .Output("output: output_types")
    .Attr("Targuments: list(type) >= 1")
    .Attr("Tcaptured: list(type) >= 0 = []")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("f: func")
    .Attr("max_intra_op_parallelism: int = 1")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      std::vector<PartialTensorShape> output_shapes;
      TF_RETURN_IF_ERROR(c->GetAttr("output_shapes", &output_shapes));
      DataTypeVector t_args;
      TF_RETURN_IF_ERROR(c->GetAttr("Targuments", &t_args));
      if (output_shapes.size() != c->num_outputs()) {
        return errors::InvalidArgument(
            "`output_shapes` must be the same length as `output_types` (",
            output_shapes.size(), " vs. ", c->num_outputs(), ")");
      }

      int64_t dim_zero = -1;
      for (size_t i = 0; i < t_args.size(); ++i) {
        if (c->Rank(c->input(i)) == 0) {
          return errors::InvalidArgument(
              "Arguments must have rank at least 1. Input ", i,
              " has rank of 0.");
        }
        auto dim_handle = c->Dim(c->input(i), 0);
        if (c->ValueKnown(dim_handle)) {
          if (dim_zero == -1) {
            dim_zero = c->Value(dim_handle);
          } else if (c->Value(dim_handle) != dim_zero) {
            return errors::InvalidArgument(
                "Arguments must have the same dimension 0.");
          }
        }
      }

      for (size_t i = 0; i < output_shapes.size(); ++i) {
        PartialTensorShape s({});
        s = s.Concatenate(dim_zero);
        s = s.Concatenate(output_shapes[i]);
        shape_inference::ShapeHandle output_shape_handle;

        TF_RETURN_IF_ERROR(
            c->MakeShapeFromPartialTensorShape(s, &output_shape_handle));
        c->set_output(static_cast<int>(i), output_shape_handle);
      }
      return absl::OkStatus();
    });

REGISTER_OP("WrapDatasetVariant")
    .Input("input_handle: variant")
    .Output("output_handle: variant")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("UnwrapDatasetVariant")
    .Input("input_handle: variant")
    .Output("output_handle: variant")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("AnonymousMultiDeviceIterator")
    .Output("handle: resource")
    .Output("deleter: variant")
    .Attr("devices: list(string) >= 1")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      c->set_output(1, c->Scalar());
      return absl::OkStatus();
    });

REGISTER_OP("AnonymousMultiDeviceIteratorV3")
    .Output("handle: resource")
    .Attr("devices: list(string) >= 1")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->Scalar());
      return absl::OkStatus();
    });

REGISTER_OP("MultiDeviceIterator")
    .Output("handle: resource")
    .Attr("devices: list(string) >= 1")
    .Attr("shared_name: string")
    .Attr("container: string")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("MultiDeviceIteratorInit")
    .Input("dataset: variant")
    .Input("multi_device_iterator: resource")
    .Input("max_buffer_size: int64")
    .Output("incarnation_id: int64")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("MultiDeviceIteratorGetNextFromShard")
    .Input("multi_device_iterator: resource")
    .Input("shard_num: int32")
    .Input("incarnation_id: int64")
    .Output("components: output_types")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetShapeFn(shape_inference::DatasetIteratorShape);

REGISTER_OP("MultiDeviceIteratorToStringHandle")
    .Input("multi_device_iterator: resource")
    .Output("string_handle: string")
    .SetForwardTypeFn(full_type::Encode(TFT_STRING, 0))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("MultiDeviceIteratorFromStringHandle")
    .Input("string_handle: string")
    .Output("multi_device_iterator: resource")
    .Attr("output_types: list(type) >= 0 = []")
    .Attr("output_shapes: list(shape) >= 0 = []")
    .SetForwardTypeFn(full_type::Decode(TFT_STRING, 0))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("OptionsDataset")
    .Input("input_dataset: variant")
    .Output("handle: variant")
    .Attr("serialized_options: string")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .Attr("metadata: string = ''")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("GetOptions")
    .Input("input_dataset: variant")
    .Output("serialized_options: string")
    .SetShapeFn(shape_inference::ScalarShape);

REGISTER_OP("FinalizeDataset")
    .Input("input_dataset: variant")
    .Output("handle: variant")
    .Attr("has_captured_ref: bool = false")
    .Attr("output_types: list(type) >= 1")
    .Attr("output_shapes: list(shape) >= 1")
    .SetTypeConstructor(full_type::VariadicTensorContainer(TFT_DATASET,
                                                           "output_types"))
    .SetShapeFn(shape_inference::ScalarShape);

}  // namespace tensorflow
