From 0c5bc309b5b1e4243f2b9921ac11db3007203ff6 Mon Sep 17 00:00:00 2001 From: Flax Team Date: Sat, 31 Aug 2024 09:35:11 -0700 Subject: [PATCH] Debug `ScopeParamNotFoundError` PiperOrigin-RevId: 669678371 --- flax/core/scope.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/flax/core/scope.py b/flax/core/scope.py index e056d6ddb..7811fa64d 100644 --- a/flax/core/scope.py +++ b/flax/core/scope.py @@ -967,6 +967,26 @@ def param( The parameters. Throws an error if the params exist already. """ self.reserve(name, 'params') + print( + f'Debugging param(): name={name}, scope_path={self.path_text}' + ) # New debugging statement + + # Check if 'params' collection exists and print its contents + # if it's not empty + params_collection = self._collection('params') + # Check if it's not a FrozenDict (empty) + if not isinstance(params_collection, FrozenDict): + print(f" - Contents of 'params' collection in scope {self.path_text}:") + for param_name, param_value in params_collection.items(): + print(f" - {param_name}: {param_value.shape}") + # Print the structure of `self._variables` + print(f" - Structure of '_variables' in scope {self.path_text}:") + for col_name, collection in self._variables.items(): + print(f" - {col_name}:") + if isinstance(collection, dict): + for var_name, var_value in collection.items(): + print(f" - {var_name}: {var_value.shape}") + if self.has_variable('params', name): value = self.get_variable('params', name) if unbox: