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: