From 54a205b6591aa789a44e4ca581c1fdc524ee6b74 Mon Sep 17 00:00:00 2001 From: Ivy Zheng Date: Wed, 22 Feb 2023 16:46:17 -0800 Subject: [PATCH] Explicitly mark the lazy-rng with legacy_config, and all new configs should use flax.config. PiperOrigin-RevId: 511631913 --- flax/core/scope.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/flax/core/scope.py b/flax/core/scope.py index 081b929f2..c69011556 100644 --- a/flax/core/scope.py +++ b/flax/core/scope.py @@ -23,7 +23,8 @@ Sequence, Set, Tuple, TypeVar, Union) from flax.ids import uuid -from flax import configurations as config +from flax import config as config +from flax import configurations as legacy_config # only for flax_lazy_rng from flax import errors from flax import struct from flax import traceback_util @@ -96,7 +97,7 @@ def as_jax_rng(self) -> PRNGKey: @staticmethod def create(rng: Union['LazyRng', PRNGKey], *suffix: PRNGFoldable) -> 'LazyRng': - if not config.flax_lazy_rng: + if not legacy_config.flax_lazy_rng: if isinstance(rng, LazyRng): assert not rng.suffix rng = rng.rng