Fix implicit import of flax.struct

PiperOrigin-RevId: 581881281
google · Nov 13, 2023 · 991d3b7 · 991d3b7
1 parent 2f6e5ff
commit 991d3b7
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/flax/__init__.py b/flax/__init__.py
@@ -24,6 +24,7 @@
 from . import jax_utils
 from . import linen
 from . import serialization
+from . import struct
 from . import traverse_util
 
 # DO NOT REMOVE - Marker for internal deprecated API.