[Datasets] Add `from_tf`
diff --git a/doc/requirements-doc.txt b/doc/requirements-doc.txt
index 06ba20c994ad6..9a107712401dd 100644
--- a/doc/requirements-doc.txt
+++ b/doc/requirements-doc.txt
@@ -31,6 +31,7 @@ werkzeug
 wandb
 tensorflow; sys_platform != 'darwin' or platform_machine != 'arm64'
 tensorflow-macos; sys_platform == 'darwin' and platform_machine == 'arm64'
+tensorflow-datasets
 torch
 torchvision
 transformers
diff --git a/doc/source/data/api/input_output.rst b/doc/source/data/api/input_output.rst
index ce54225461836..e1d6606e645f1 100644
--- a/doc/source/data/api/input_output.rst
+++ b/doc/source/data/api/input_output.rst
@@ -147,6 +147,11 @@ HuggingFace
 
 .. autofunction:: ray.data.from_huggingface
 
+TensorFlow
+----------
+
+.. autofunction:: ray.data.from_tf
+
 
 .. _data_source_api:
 
@@ -195,9 +200,6 @@ Built-in Datasources
 .. autoclass:: ray.data.datasource.RangeDatasource
     :members:
 
-.. autoclass:: ray.data.datasource.SimpleTensorFlowDatasource
-    :members:
-
 .. autoclass:: ray.data.datasource.TFRecordDatasource
     :members:
 
diff --git a/doc/source/data/creating-datasets.rst b/doc/source/data/creating-datasets.rst
index 06c02a418b55c..cd8d2785be8bd 100644
--- a/doc/source/data/creating-datasets.rst
+++ b/doc/source/data/creating-datasets.rst
@@ -521,29 +521,22 @@ From Torch and TensorFlow
 .. tabbed:: TensorFlow
 
     If you already have a TensorFlow dataset available, you can create a Ray Dataset
-    using :py:class:`SimpleTensorFlowDatasource`.
+    using :class:`~ray.data.from_tf`.
 
     .. warning::
-        :py:class:`SimpleTensorFlowDatasource` doesn't support parallel reads. You
-        should only use this datasource for small datasets like MNIST or CIFAR.
+        :class:`~ray.data.from_tf` doesn't support parallel reads. You
+        should only use this function with small datasets like MNIST or CIFAR.
 
     .. code-block:: python
 
-        import ray.data
-        from ray.data.datasource import SimpleTensorFlowDatasource
+        import ray
         import tensorflow_datasets as tfds
 
-        def dataset_factory():
-            return tfds.load("cifar10", split=["train"], as_supervised=True)[0]
-
-        dataset = ray.data.read_datasource(
-            SimpleTensorFlowDatasource(),
-            parallelism=1,
-            dataset_factory=dataset_factory
-        )
-        features, label = dataset.take(1)[0]
-        features.shape  # TensorShape([32, 32, 3])
-        label  # <tf.Tensor: shape=(), dtype=int64, numpy=7>
+        dataset, _ = tfds.load("cifar10", split=["train", "test"])
+        dataset = ray.data.from_tf(dataset)
+
+        dataset
+        # -> Dataset(num_blocks=200, num_rows=50000, schema={id: binary, image: ArrowTensorType(shape=(32, 32, 3), dtype=uint8), label: int64})
 
 .. _dataset_from_huggingface:
 
diff --git a/python/ray/data/__init__.py b/python/ray/data/__init__.py
index 527711745f652..1208672d71721 100644
--- a/python/ray/data/__init__.py
+++ b/python/ray/data/__init__.py
@@ -22,6 +22,7 @@
     from_pandas,
     from_pandas_refs,
     from_spark,
+    from_tf,
     from_torch,
     range,
     range_arrow,
@@ -65,6 +66,7 @@
     "from_pandas",
     "from_pandas_refs",
     "from_spark",
+    "from_tf",
     "from_torch",
     "from_huggingface",
     "range",
diff --git a/python/ray/data/datasource/__init__.py b/python/ray/data/datasource/__init__.py
index 68eb7c4ceff61..7f0c61ce440a2 100644
--- a/python/ray/data/datasource/__init__.py
+++ b/python/ray/data/datasource/__init__.py
@@ -37,7 +37,6 @@
     Partitioning,
 )
 from ray.data.datasource.tfrecords_datasource import TFRecordDatasource
-from ray.data.datasource.tensorflow_datasource import SimpleTensorFlowDatasource
 from ray.data.datasource.text_datasource import TextDatasource
 
 __all__ = [
@@ -69,7 +68,6 @@
     "RangeDatasource",
     "ReadTask",
     "Reader",
-    "SimpleTensorFlowDatasource",
     "TextDatasource",
     "TFRecordDatasource",
     "WriteResult",
diff --git a/python/ray/data/datasource/tensorflow_datasource.py b/python/ray/data/datasource/tensorflow_datasource.py
deleted file mode 100644
index 7c7e6b4260534..0000000000000
--- a/python/ray/data/datasource/tensorflow_datasource.py
+++ /dev/null
@@ -1,73 +0,0 @@
-import logging
-from typing import TYPE_CHECKING, Callable, Iterator, List
-
-from ray.data.block import Block, BlockMetadata, T
-from ray.data.datasource import Datasource, ReadTask
-from ray.util.annotations import PublicAPI
-
-if TYPE_CHECKING:
-    import tensorflow as tf
-
-logger = logging.getLogger(__name__)
-
-
-@PublicAPI
-class SimpleTensorFlowDatasource(Datasource[T]):
-    """A datasource that lets you use TensorFlow datasets with Ray Data.
-
-    .. warning::
-        ``SimpleTensorFlowDataset`` doesn't support parallel reads. You should only use
-        this datasource for small datasets like MNIST or CIFAR.
-
-    Example:
-        >>> import ray.data
-        >>> from ray.data.datasource import SimpleTensorFlowDatasource
-        >>> import tensorflow_datasets as tfds  # doctest: +SKIP
-        >>>
-        >>> def dataset_factory():
-        ...     return tfds.load("cifar10", split=["train"], as_supervised=True)[0]
-        ...
-        >>> dataset = ray.data.read_datasource(  # doctest: +SKIP
-        ...     SimpleTensorFlowDatasource(),
-        ...     parallelism=1,
-        ...     dataset_factory=dataset_factory
-        ... )
-        >>> features, label = dataset.take(1)[0]  # doctest: +SKIP
-        >>> features.shape  # doctest: +SKIP
-        TensorShape([32, 32, 3])
-        >>> label  # doctest: +SKIP
-        <tf.Tensor: shape=(), dtype=int64, numpy=7>
-    """
-
-    def prepare_read(
-        self,
-        parallelism: int,
-        dataset_factory: Callable[[], "tf.data.Dataset"],
-    ) -> List[ReadTask]:
-        """Return a read task that loads a TensorFlow dataset.
-
-        Arguments:
-            parallelism: This argument isn't used.
-            dataset_factory: A no-argument function that returns the TensorFlow dataset
-                to be read.
-        """
-        if parallelism > 1:
-            logger.warn(
-                "`SimpleTensorFlowDatasource` doesn't support parallel reads. The "
-                "`parallelism` argument will be ignored."
-            )
-
-        def read_fn() -> Iterator[Block]:
-            # Load the entire dataset into memory.
-            block = list(dataset_factory())
-            # Store the data in a single block.
-            yield block
-
-        metadata = BlockMetadata(
-            num_rows=None,
-            size_bytes=None,
-            schema=None,
-            input_files=None,
-            exec_stats=None,
-        )
-        return [ReadTask(read_fn, metadata)]
diff --git a/python/ray/data/read_api.py b/python/ray/data/read_api.py
index dab8e99c2bb76..11a7711eca8e8 100644
--- a/python/ray/data/read_api.py
+++ b/python/ray/data/read_api.py
@@ -59,6 +59,7 @@
     import pandas
     import pyarrow
     import pyspark
+    import tensorflow as tf
     import torch
 
 
@@ -1341,10 +1342,10 @@ def convert(ds: "datasets.Dataset") -> Dataset[ArrowRow]:
 
 
 @PublicAPI
-def from_torch(
-    dataset: "torch.utils.data.Dataset",
+def from_tf(
+    dataset: "tf.data.Dataset",
 ) -> Dataset:
-    """Create a dataset from a Torch dataset.
+    """Create a dataset from a TensorFlow dataset.
 
     This function is inefficient. Use it to read small datasets or prototype.
 
@@ -1357,6 +1358,55 @@ def from_torch(
         This function isn't paralellized. It loads the entire dataset into the head
         node's memory before moving the data to the distributed object store.
 
+    Examples:
+        >>> import ray
+        >>> import tensorflow_datasets as tfds
+        >>> dataset, _ = tfds.load('cifar10', split=["train", "test"])  # doctest: +SKIP
+        >>> dataset = ray.data.from_tf(dataset)  # doctest: +SKIP
+        >>> dataset  # doctest: +SKIP
+        Dataset(num_blocks=200, num_rows=50000, schema={id: binary, image: ArrowTensorType(shape=(32, 32, 3), dtype=uint8), label: int64})
+        >>> dataset.take(1)  # doctest: +SKIP
+        [{'id': b'train_16399', 'image': array([[[143,  96,  70],
+        [141,  96,  72],
+        [135,  93,  72],
+        ...,
+        [ 96,  37,  19],
+        [105,  42,  18],
+        [104,  38,  20]],
+
+       ...,
+
+       [[195, 161, 126],
+        [187, 153, 123],
+        [186, 151, 128],
+        ...,
+        [212, 177, 147],
+        [219, 185, 155],
+        [221, 187, 157]]], dtype=uint8), 'label': 7}]
+
+    Args:
+        dataset: A TensorFlow dataset.
+
+    Returns:
+        A :class:`Dataset` that contains the samples stored in the TensorFlow dataset.
+    """  # noqa: E501
+    # FIXME: `as_numpy_iterator` errors if `dataset` contains ragged tensors.
+    return from_items(list(dataset.as_numpy_iterator()))
+
+
+@PublicAPI
+def from_torch(
+    dataset: "torch.utils.data.Dataset",
+) -> Dataset:
+    """Create a dataset from a Torch dataset.
+    This function is inefficient. Use it to read small datasets or prototype.
+    .. warning::
+        If your dataset is large, this function may execute slowly or raise an
+        out-of-memory error. To avoid issues, read the underyling data with a function
+        like :meth:`~ray.data.read_images`.
+    .. note::
+        This function isn't paralellized. It loads the entire dataset into the head
+        node's memory before moving the data to the distributed object store.
     Examples:
         >>> import ray
         >>> from torchvision import datasets
@@ -1366,10 +1416,8 @@ def from_torch(
         Dataset(num_blocks=200, num_rows=60000, schema=<class 'tuple'>)
         >>> dataset.take(1)  # doctest: +SKIP
         [(<PIL.Image.Image image mode=L size=28x28 at 0x...>, 5)]
-
     Args:
         dataset: A Torch dataset.
-
     Returns:
         A :class:`Dataset` that contains the samples stored in the Torch dataset.
     """
diff --git a/python/ray/data/tests/test_dataset_formats.py b/python/ray/data/tests/test_dataset_formats.py
index 5caffce4dd69d..d6af782f41234 100644
--- a/python/ray/data/tests/test_dataset_formats.py
+++ b/python/ray/data/tests/test_dataset_formats.py
@@ -17,7 +17,6 @@
 from ray.data.datasource import (
     Datasource,
     DummyOutputDatasource,
-    SimpleTensorFlowDatasource,
     WriteResult,
 )
 
@@ -192,23 +191,18 @@ def test_write_datasource(ray_start_regular_shared, pipelined):
     assert ray.get(output.data_sink.get_rows_written.remote()) == 10
 
 
-def test_tensorflow_datasource(ray_start_regular_shared):
+def test_from_tf(ray_start_regular_shared):
     import tensorflow as tf
     import tensorflow_datasets as tfds
 
     tf_dataset = tfds.load("mnist", split=["train"], as_supervised=True)[0]
+    tf_dataset = tf_dataset.take(8)  # Use subset to make test run faster.
 
-    def dataset_factory():
-        return tfds.load("mnist", split=["train"], as_supervised=True)[0]
-
-    ray_dataset = ray.data.read_datasource(
-        SimpleTensorFlowDatasource(), parallelism=1, dataset_factory=dataset_factory
-    ).fully_executed()
-
-    assert ray_dataset.num_blocks() == 1
+    ray_dataset = ray.data.from_tf(tf_dataset)
 
     actual_data = ray_dataset.take_all()
     expected_data = list(tf_dataset)
+    assert len(actual_data) == len(expected_data)
     for (expected_features, expected_label), (actual_features, actual_label) in zip(
         expected_data, actual_data
     ):

Signed-off-by: Balaji <balaji@anyscale.com>

**Blocked on**
- [x] #29590 

<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this solves. -->

Parity with `from_huggingface`. Also, `SimpleTensorflowDatasource` is clunky. For more motivation, read #29430.

## Related issue number

<!-- For example: "Closes #1234" -->

Towards #29430. See also #29589.

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR.
- [x] I've run `scripts/format.sh` to lint the changes in this PR.
- [x] I've included any doc changes needed for https://docs.ray.io/en/master/.
- [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(
