diff --git a/tensorflow/core/ops/dataset_ops.cc b/tensorflow/core/ops/dataset_ops.cc index 6a633fb679d..4f750cc938d 100644 --- a/tensorflow/core/ops/dataset_ops.cc +++ b/tensorflow/core/ops/dataset_ops.cc @@ -415,6 +415,11 @@ REGISTER_OP("AnonymousSeedGenerator") return Status::OK(); }); +REGISTER_OP("DatasetCardinality") + .Input("input_dataset: variant") + .Output("cardinality: int64") + .SetShapeFn(shape_inference::ScalarShape); + REGISTER_OP("DeleteSeedGenerator") .Input("handle: resource") .Input("deleter: variant") diff --git a/tensorflow/core/ops/experimental_dataset_ops.cc b/tensorflow/core/ops/experimental_dataset_ops.cc index 18b35d3a172..ac36540e29a 100644 --- a/tensorflow/core/ops/experimental_dataset_ops.cc +++ b/tensorflow/core/ops/experimental_dataset_ops.cc @@ -229,11 +229,6 @@ REGISTER_OP("ExperimentalCSVDataset") return shape_inference::ScalarShape(c); }); -REGISTER_OP("DatasetCardinality") - .Input("input_dataset: variant") - .Output("cardinality: int64") - .SetShapeFn(shape_inference::ScalarShape); - REGISTER_OP("ExperimentalDatasetCardinality") .Input("input_dataset: variant") .Output("cardinality: int64") diff --git a/tensorflow/python/data/experimental/ops/cardinality.py b/tensorflow/python/data/experimental/ops/cardinality.py index f1b8908fa68..3505d1fc87e 100644 --- a/tensorflow/python/data/experimental/ops/cardinality.py +++ b/tensorflow/python/data/experimental/ops/cardinality.py @@ -20,6 +20,7 @@ from __future__ import print_function from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.ops import gen_dataset_ops from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops from tensorflow.python.util.tf_export import tf_export @@ -64,7 +65,7 @@ def cardinality(dataset): constant `INFINITE_CARDINALITY` and `UNKNOWN_CARDINALITY` respectively. """ - return ged_ops.dataset_cardinality(dataset._variant_tensor) # pylint: disable=protected-access + return gen_dataset_ops.dataset_cardinality(dataset._variant_tensor) # pylint: disable=protected-access @tf_export("data.experimental.assert_cardinality") diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py index 372c19855af..1b2b7c779b1 100644 --- a/tensorflow/python/data/ops/dataset_ops.py +++ b/tensorflow/python/data/ops/dataset_ops.py @@ -2157,7 +2157,7 @@ name=None)) named constants `tf.data.INFINITE_CARDINALITY` and `tf.data.UNKNOWN_CARDINALITY` respectively. """ - return ged_ops.dataset_cardinality(self._variant_tensor) + return gen_dataset_ops.dataset_cardinality(self._variant_tensor) @tf_export(v1=["data.Dataset"])