Deprecation Warning
Hello,
On page: [TensorFlow Datasets](https://www.tensorflow.org/datasets/overview), the following line of code under the title [```As batched tf.Tensor (batch_size=-1)```](https://www.tensorflow.org/datasets/overview#as_batched_tftensor_batch_size-1) will generate a **Deprecation Warning**.
```python
image, label = tfds.as_numpy(tfds.load(
    'mnist',
    split='test',
    batch_size=-1,
    as_supervised=True,
))

print(type(image), image.shape)
```

The warning message is:
```
WARNING:tensorflow:From e:\\python39\\lib\\site-packages\\tensorflow_datasets\\core\\dataset_builder.py:622: get_single_element (from tensorflow.python.data.experimental.ops.get_single_element) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Dataset.get_single_element()`.
```
Thank you,
diff --git a/tensorflow_datasets/core/dataset_builder.py b/tensorflow_datasets/core/dataset_builder.py
index f9f313e62..4f50e9c58 100644
--- a/tensorflow_datasets/core/dataset_builder.py
+++ b/tensorflow_datasets/core/dataset_builder.py
@@ -619,7 +619,7 @@ class DatasetBuilder(registered.RegisteredDataset):
     # allow the user to overwritte it.
 
     if wants_full_dataset:
-      return tf.data.experimental.get_single_element(ds)
+      return tf_compat.get_single_element(ds)
     return ds
 
   def _should_cache_ds(self, split, shuffle_files, read_config):
diff --git a/tensorflow_datasets/core/tf_compat.py b/tensorflow_datasets/core/tf_compat.py
index 8ed433f1f..d489f4b71 100644
--- a/tensorflow_datasets/core/tf_compat.py
+++ b/tensorflow_datasets/core/tf_compat.py
@@ -71,6 +71,15 @@ def is_dataset(ds):
   return isinstance(ds, (tf.data.Dataset, tf.compat.v1.data.Dataset))
 
 
+def get_single_element(ds):
+  """Calls `tf.data.Dataset.get_single_element`."""
+  import tensorflow as tf  # pylint: disable=import-outside-toplevel
+  if hasattr(ds, "get_single_element"):  # tf 2.6 and above
+    return ds.get_single_element()
+  else:
+    return tf.data.experimental.get_single_element(ds)
+
+
 def _make_pathlike_fn(fn, nb_path_arg=1):
   """Wrap the function in a PathLike-compatible function."""
 
['tensorflow_datasets/core/tf_compat.py', 'tensorflow_datasets/core/dataset_builder.py']