From 5e5d811af21a23266c9bb1a7d9e5a80bc6bde218 Mon Sep 17 00:00:00 2001
From: Etienne Pot <epot@google.com>
Date: Mon, 20 Jan 2025 01:54:55 -0800
Subject: [PATCH] Add nnx to lazy imports

PiperOrigin-RevId: 717447982
---
 etils/ecolab/lazy_imports.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/etils/ecolab/lazy_imports.py b/etils/ecolab/lazy_imports.py
index e8338861..a13d1c8a 100644
--- a/etils/ecolab/lazy_imports.py
+++ b/etils/ecolab/lazy_imports.py
@@ -154,6 +154,7 @@ def print_current_imports() -> None:
   import flask
   import flax
   from flax import linen as nn
+  from flax import nnx
   import functorch
   import gin
   import grain.python as grain