Repository: google/flax

Issue:
nn.Module.tabulate double counting parameters
### System information
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: 0.6.6
- Python version: 3.10
- 
### Problem you have encountered:
Module summary row in tabulate is counted along with per layer/inner module param counts when computing total

### What you expected to happen:
Sample model should show a total of 50 parameters, not 100.


```
class Foo(nn.Module):
    @nn.compact
    def __call__(self, x):
        h = nn.Dense(4)(x)
        return nn.Dense(2)(h)

x = jnp.ones((16, 9))

print(Foo().tabulate(jax.random.PRNGKey(0), x)
```


```
                                 Foo Summary                                 
┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ path    ┃ module ┃ inputs        ┃ outputs       ┃ params                 ┃
┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━┩
│         │ Foo    │ float32[16,9] │ float32[16,2] │ Dense_0:               │
│         │        │               │               │   bias: float32[4]     │
│         │        │               │               │   kernel: float32[9,4] │
│         │        │               │               │ Dense_1:               │
│         │        │               │               │   bias: float32[2]     │
│         │        │               │               │   kernel: float32[4,2] │
│         │        │               │               │                        │
│         │        │               │               │ 50 (200 B)             │
├─────────┼────────┼───────────────┼───────────────┼────────────────────────┤
│ Dense_0 │ Dense  │ float32[16,9] │ float32[16,4] │ bias: float32[4]       │
│         │        │               │               │ kernel: float32[9,4]   │
│         │        │               │               │                        │
│         │        │               │               │ 40 (160 B)             │
├─────────┼────────┼───────────────┼───────────────┼────────────────────────┤
│ Dense_1 │ Dense  │ float32[16,4] │ float32[16,2] │ bias: float32[2]       │
│         │        │               │               │ kernel: float32[4,2]   │
│         │        │               │               │                        │
│         │        │               │               │ 10 (40 B)              │
├─────────┼────────┼───────────────┼───────────────┼────────────────────────┤
│         │        │               │         Total │ 100 (400 B)            │
└─────────┴────────┴───────────────┴───────────────┴────────────────────────┘
                                                                             
                        Total Parameters: 100 (400 B)          
```



Discussion:
Thanks @oars for reporting this! I've sent a patch, should be fixed soon. 
Will make a release once its merged.

Resolving patch:
diff --git a/flax/linen/summary.py b/flax/linen/summary.py
--- a/flax/linen/summary.py
+++ b/flax/linen/summary.py
@@ -17,6 +17,7 @@
 import dataclasses
 import io
 from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union
+from flax.core import unfreeze
 
 import flax.linen.module as module_lib
 from flax.core import meta
@@ -310,7 +311,7 @@ def _get_path_variables(path: Tuple[str, ...], variables: FrozenVariableDict) ->
   path_variables = {}
 
   for collection in variables:
-    collection_variables = jax.tree_util.tree_map(lambda x: x, variables[collection]) # make a deep copy
+    collection_variables = variables[collection]
     for name in path:
       if name not in collection_variables:
         collection_variables = None
@@ -318,7 +319,7 @@ def _get_path_variables(path: Tuple[str, ...], variables: FrozenVariableDict) ->
       collection_variables = collection_variables[name]
 
     if collection_variables is not None:
-      path_variables[collection] = collection_variables
+      path_variables[collection] = unfreeze(collection_variables)
 
   return path_variables
 


Test patch:
diff --git a/tests/linen/summary_test.py b/tests/linen/summary_test.py
--- a/tests/linen/summary_test.py
+++ b/tests/linen/summary_test.py
@@ -571,6 +571,18 @@ def __call__(self):
     self.assertIn('baz', lines[7])
     self.assertIn('qux', lines[8])
 
+  def test_tabulate_param_count(self):
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        h = nn.Dense(4)(x)
+        return nn.Dense(2)(h)
+
+    x = jnp.ones((16, 9))
+    rep = Foo().tabulate(jax.random.PRNGKey(0), x, console_kwargs=CONSOLE_TEST_KWARGS)
+    lines = rep.splitlines()
+    self.assertIn('Total Parameters: 50', lines[-2])
+
 
 if __name__ == '__main__':
   absltest.main()