From 11bd2b0a9ca90edcfbba9d5fb9c5eacaa41fcf79 Mon Sep 17 00:00:00 2001
From: Grain Team <grain-dev@google.com>
Date: Tue, 2 Jun 2026 08:31:10 -0700
Subject: [PATCH] Enable lazy imports for stats and stats_utils targets to
 improve startup time.

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

diff --git a/grain/_src/python/dataset/BUILD b/grain/_src/python/dataset/BUILD
index acbafa81a..ff44cb644 100644
--- a/grain/_src/python/dataset/BUILD
+++ b/grain/_src/python/dataset/BUILD
@@ -95,6 +95,7 @@ py_test(
 py_library(
     name = "stats",
     srcs = ["stats.py"],
+    lazy_imports = True,
     srcs_version = "PY3",
     deps = [
         ":base",
@@ -111,6 +112,7 @@ py_library(
 py_library(
     name = "stats_utils",
     srcs = ["stats_utils.py"],
+    lazy_imports = True,
     srcs_version = "PY3",
     deps = [
         "//grain/_src/core:tree_lib",