From 73e16900cd3dd4f9ce1b474066d9182c71eb9e70 Mon Sep 17 00:00:00 2001
From: Grain Team <grain-dev@google.com>
Date: Tue, 7 Apr 2026 12:40:19 -0700
Subject: [PATCH] Enable lazy imports for
 //third_party/py/grain/_src/python/dataset:dataset.

PiperOrigin-RevId: 896038784
---
 grain/_src/python/dataset/BUILD | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/grain/_src/python/dataset/BUILD b/grain/_src/python/dataset/BUILD
index 8fcd0bbcd..718abccc5 100644
--- a/grain/_src/python/dataset/BUILD
+++ b/grain/_src/python/dataset/BUILD
@@ -37,11 +37,12 @@ py_library(
         "dataset.py",
         "//grain/_src/python/dataset/transformations:core_transformations",
     ],
-    srcs_version = "PY3",
     data = select({
         "@platforms//os:windows": ["//grain/_src/python/experimental/index_shuffle/python:index_shuffle_module.pyd"],
         "//conditions:default": ["//grain/_src/python/experimental/index_shuffle/python:index_shuffle_module.so"],
     }),
+    lazy_imports = True,
+    srcs_version = "PY3",
     deps = [
         ":base",
         ":stats",