diff --git a/kauldron/data/__init__.py b/kauldron/data/__init__.py
index a02d1024..280f563c 100644
--- a/kauldron/data/__init__.py
+++ b/kauldron/data/__init__.py
@@ -37,14 +37,20 @@
 
   from kauldron.data.utils import BatchSize
 
-# ****************************************************************************
-# DO NOT ADD preprocessing ops here. Instead, add them to `kd.contrib.data`
-# ****************************************************************************
-
-# TODO(epot): Should migrate all users to use explicitly `kd.data.tf`
-from kauldron.data.transforms.base import Elements
-from kauldron.data.transforms.base import ElementWiseTransform
-from kauldron.data.transforms.base import TreeFlattenWithPath
-from kauldron.data.transforms.map_transforms import Gather
-from kauldron.data.transforms.map_transforms import Rearrange
-from kauldron.data.transforms.map_transforms import ValueRange
+  from kauldron.data.py.data_sources import DataSource
+  from kauldron.data.py.data_sources import Json
+  from kauldron.data.py.data_sources import HuggingFace
+  from kauldron.data.py.data_sources import Tfds
+  from kauldron.data.py.mixtures import Mix
+
+  # ****************************************************************************
+  # DO NOT ADD preprocessing ops here. Instead, add them to `kd.contrib.data`
+  # ****************************************************************************
+
+  # TODO(epot): Should migrate all users to use explicitly `kd.data.tf`
+  from kauldron.data.transforms.base import Elements
+  from kauldron.data.transforms.base import ElementWiseTransform
+  from kauldron.data.transforms.base import TreeFlattenWithPath
+  from kauldron.data.transforms.map_transforms import Gather
+  from kauldron.data.transforms.map_transforms import Rearrange
+  from kauldron.data.transforms.map_transforms import ValueRange