Make a number of functions visible.

PiperOrigin-RevId: 576175489
google · Oct 24, 2023 · fd8fd76 · fd8fd76
1 parent abe205b
commit fd8fd76
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/flax/linen/normalization.py b/flax/linen/normalization.py
@@ -32,8 +32,13 @@
 Shape = Tuple[int, ...]
 Dtype = Any  # this could be a real type?
 Axes = Union[int, Sequence[int]]
+
+field = dataclasses.field
+canonicalize_dtype = dtypes.canonicalize_dtype
 compact = module.compact
 Module = module.Module
+merge_param = module.merge_param
+map_variables = transforms.map_variables
 
 
 def _canonicalize_axes(rank: int, axes: Axes) -> Tuple[int, ...]: