From 6cb74f68a4d206e7855e52179af13e4dac617f79 Mon Sep 17 00:00:00 2001
From: Ihor Indyk <iindyk@google.com>
Date: Mon, 22 Dec 2025 14:56:47 -0800
Subject: [PATCH] Internal change.

PiperOrigin-RevId: 847894347
---
 grain/__init__.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/grain/__init__.py b/grain/__init__.py
index 4ebc6192a..90bd1e61b 100644
--- a/grain/__init__.py
+++ b/grain/__init__.py
@@ -18,6 +18,7 @@
 # pylint: disable=unused-import
 # pylint: disable=g-multiple-import
 # pylint: disable=g-import-not-at-top
+# pylint: disable=g-bad-import-order
 
 # We import all public modules here to enable the use of `grain.foo.Bar`
 # instead of forcing users to write `from grain import foo as grain_foo`.
@@ -31,6 +32,11 @@
     sources,
     transforms,
 )
+from grain._src.python.dataset.dataset import (
+    DatasetIterator,
+    IterDataset,
+    MapDataset,
+)
 
 from grain._src.core.config import config
 from grain._src.core.version import __version__, __version_info__
@@ -38,11 +44,7 @@
     DataLoader,
     DataLoaderIterator,
 )
-from grain._src.python.dataset.dataset import (
-    DatasetIterator,
-    IterDataset,
-    MapDataset,
-)
+
 from grain._src.python.load import load
 from grain._src.python.options import ReadOptions, MultiprocessingOptions
 from grain._src.python.record import Record, RecordMetadata