Add nnx to lazy imports

PiperOrigin-RevId: 717447982
google · Jan 20, 2025 · 5e5d811 · 5e5d811
1 parent da76653
commit 5e5d811
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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