Migrate tf.data.experimental.cardinality to tf.data.Dataset.cardinality.

PiperOrigin-RevId: 733923403

Author: yangustc07

tf_keras/engine/data_adapter.py

@@ -794,8 +794,7 @@ def should_recreate_iterator(self):
         # each epoch.
         return (
             self._user_steps is None
-            or tf.data.experimental.cardinality(self._dataset).numpy()
-            == self._user_steps
+            or self._dataset.cardinality().numpy() == self._user_steps
         )
 
     def _validate_args(self, y, sample_weights, steps, pss_evaluation_shards):
@@ -819,8 +818,8 @@ def _validate_args(self, y, sample_weights, steps, pss_evaluation_shards):
                         "specify the number of steps to run."
                     )
             else:
-                size = tf.data.experimental.cardinality(self._dataset).numpy()
-                if size == tf.data.experimental.INFINITE_CARDINALITY:
+                size = self._dataset.cardinality().numpy()
+                if size == tf.data.INFINITE_CARDINALITY:
                     if pss_evaluation_shards:
                         raise ValueError(
                             "When performing exact evaluation, the dataset "
@@ -1481,8 +1480,8 @@ def _infer_steps(self, steps, dataset):
         if not isinstance(dataset, tf.data.Dataset):
             return None
 
-        size = tf.data.experimental.cardinality(dataset)
-        if size == tf.data.experimental.INFINITE_CARDINALITY and steps is None:
+        size = dataset.cardinality()
+        if size == tf.data.INFINITE_CARDINALITY and steps is None:
             raise ValueError(
                 "When passing an infinitely repeating dataset, please specify "
                 "a `steps_per_epoch` value so that epoch level "

tf_keras/engine/data_adapter_test.py

@@ -1198,8 +1198,8 @@ def test_unknown_cardinality_dataset_with_steps_per_epoch(self):
         ds = tf.data.Dataset.from_tensor_slices([0, 1, 2, 3, 4, 5, 6])
         filtered_ds = ds.filter(lambda x: x < 4)
         self.assertEqual(
-            tf.data.experimental.cardinality(filtered_ds).numpy(),
-            tf.data.experimental.UNKNOWN_CARDINALITY,
+            filtered_ds.cardinality().numpy(),
+            tf.data.UNKNOWN_CARDINALITY,
         )
 
         # User can choose to only partially consume `Dataset`.
@@ -1221,8 +1221,8 @@ def test_unknown_cardinality_dataset_without_steps_per_epoch(self):
         ds = tf.data.Dataset.from_tensor_slices([0, 1, 2, 3, 4, 5, 6])
         filtered_ds = ds.filter(lambda x: x < 4)
         self.assertEqual(
-            tf.data.experimental.cardinality(filtered_ds).numpy(),
-            tf.data.experimental.UNKNOWN_CARDINALITY,
+            filtered_ds.cardinality().numpy(),
+            tf.data.UNKNOWN_CARDINALITY,
         )
 
         data_handler = data_adapter.DataHandler(