From 31def0cc68754bb8e09824ba4bdad8a4e1336a99 Mon Sep 17 00:00:00 2001
From: Grain Team <grain-dev@google.com>
Date: Tue, 19 May 2026 03:06:19 -0700
Subject: [PATCH] Enable lazy imports for more JAX internal and experimental
 targets.

PiperOrigin-RevId: 917723631
---
 grain/_src/python/checkpoint/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/grain/_src/python/checkpoint/BUILD b/grain/_src/python/checkpoint/BUILD
index a160c9c46..200572c93 100644
--- a/grain/_src/python/checkpoint/BUILD
+++ b/grain/_src/python/checkpoint/BUILD
@@ -15,6 +15,7 @@ py_library(
 py_library(
     name = "handler",
     srcs = ["handler.py"],
+    lazy_imports = True,
     srcs_version = "PY3",
     deps = [
         "//grain/_src/core:sharding",