Repository: google/flax

Issue:
`nnx.tabulate` does not support `nnx.Module` with multiple inputs/outputs
Hi! 👋
First of all, thank you for this great project!

I found a small issue with the `nnx.tabulate` function.
Currently, it only supports nnx.Module instances that **take one input and return one output**.
 
You can check with following simple examples.

```
class CustomMLP(nnx.Module):
    def __init__(self):
        self.weight = nnx.Param(jnp.ones((4, 8)))
        self.bias = nnx.Param(jnp.ones(8))
    
    def __call__(self, x):
        y = x @ self.weight
        y += self.bias
        return y, x

cmlp = CustomMLP()
x = jnp.ones((1, 4))
print(nnx.tabulate(cmlp, x))
```

```
class CustomMLP(nnx.Module):
    def __init__(self):
        self.weight = nnx.Param(jnp.ones((4, 8)))
        self.bias = nnx.Param(jnp.ones(8))
    
    def __call__(self, x, x2):
        y = x @ self.weight
        y += self.bias
        y += x2
        return y

cmlp = CustomMLP()
x = jnp.ones((1, 4))
x2 = jnp.ones((1, 8))
print(nnx.tabulate(cmlp, x, x2))
```

Discussion:

Solution patch:
diff --git a/README.md b/README.md
index 42321ead9..e8de95600 100644
--- a/README.md
+++ b/README.md
@@ -157,7 +157,7 @@ To cite this repository:
   author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},
   title = {{F}lax: A neural network library and ecosystem for {JAX}},
   url = {http://github.com/google/flax},
-  version = {0.10.3},
+  version = {0.10.4},
   year = {2024},
 }
 ```
diff --git a/flax/nnx/summary.py b/flax/nnx/summary.py
index 7f1aed533..7396b2edc 100644
--- a/flax/nnx/summary.py
+++ b/flax/nnx/summary.py
@@ -519,7 +519,12 @@ def _unflatten_to_simple_structure(xs: list[tuple[tuple[tp.Any, ...], tp.Any]]):
           else:
             cursor[key] = {}
       cursor = cursor[key]
-    cursor[path[-1]] = value
+    if isinstance(cursor, list):
+      assert path[-1] == len(cursor)
+      cursor.append(value)
+    else:
+      assert isinstance(cursor, dict)
+      cursor[path[-1]] = value
   return result
 
 def _as_yaml_str(value) -> str:
diff --git a/flax/version.py b/flax/version.py
index 2f12c4c34..ac1d8a158 100644
--- a/flax/version.py
+++ b/flax/version.py
@@ -13,4 +13,4 @@
 # limitations under the License.
 
 """Current Flax version at head on Github."""
-__version__ = '0.10.3'
+__version__ = '0.10.4'
diff --git a/uv.lock b/uv.lock
index 6d1e26f27..cd6a412c3 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2472,6 +2472,8 @@ version = "6.0.0"
 source = { registry = "https://pypi.org/simple" }
 sdist = { url = "https://files.pythonhosted.org/packages/18/c7/8c6872f7372eb6a6b2e4708b88419fb46b857f7a2e1892966b851cc79fc9/psutil-6.0.0.tar.gz", hash = "sha256:8faae4f310b6d969fa26ca0545338b21f73c6b15db7c4a8d934a5482faa818f2", size = 508067 }
 wheels = [
+    { url = "https://files.pythonhosted.org/packages/c5/66/78c9c3020f573c58101dc43a44f6855d01bbbd747e24da2f0c4491200ea3/psutil-6.0.0-cp27-none-win32.whl", hash = "sha256:02b69001f44cc73c1c5279d02b30a817e339ceb258ad75997325e0e6169d8b35", size = 249766 },
+    { url = "https://files.pythonhosted.org/packages/e1/3f/2403aa9558bea4d3854b0e5e567bc3dd8e9fbc1fc4453c0aa9aafeb75467/psutil-6.0.0-cp27-none-win_amd64.whl", hash = "sha256:21f1fb635deccd510f69f485b87433460a603919b45e2a324ad65b0cc74f8fb1", size = 253024 },
     { url = "https://files.pythonhosted.org/packages/0b/37/f8da2fbd29690b3557cca414c1949f92162981920699cd62095a984983bf/psutil-6.0.0-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:c588a7e9b1173b6e866756dde596fd4cad94f9399daf99ad8c3258b3cb2b47a0", size = 250961 },
     { url = "https://files.pythonhosted.org/packages/35/56/72f86175e81c656a01c4401cd3b1c923f891b31fbcebe98985894176d7c9/psutil-6.0.0-cp36-abi3-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6ed2440ada7ef7d0d608f20ad89a04ec47d2d3ab7190896cd62ca5fc4fe08bf0", size = 287478 },
     { url = "https://files.pythonhosted.org/packages/19/74/f59e7e0d392bc1070e9a70e2f9190d652487ac115bb16e2eff6b22ad1d24/psutil-6.0.0-cp36-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5fd9a97c8e94059b0ef54a7d4baf13b405011176c3b6ff257c247cae0d560ecd", size = 290455 },


Test patch:
diff --git a/tests/nnx/summary_test.py b/tests/nnx/summary_test.py
index 94059c212..2555fde22 100644
--- a/tests/nnx/summary_test.py
+++ b/tests/nnx/summary_test.py
@@ -72,3 +72,31 @@ def __call__(self, x):
     self.assertIn('5,790 (23.2 KB)', table_repr[34])
     self.assertIn('2 (12 B)', table_repr[34])
     self.assertIn('Total Parameters: 6,068 (24.3 KB)', table_repr[37])
+
+  def test_multiple_inputs_and_outputs(self):
+    class CustomMLP(nnx.Module):
+      def __init__(self):
+        self.weight = nnx.Param(jnp.ones((4, 8)))
+        self.bias = nnx.Param(jnp.ones(8))
+
+      def __call__(self, x, x2):
+        y = x @ self.weight
+        y += self.bias[None]
+        y += x2
+        return x, y, 2 * y
+
+    cmlp = CustomMLP()
+    x = jnp.ones((1, 4))
+    x2 = jnp.ones((1, 8))
+    table_repr = nnx.tabulate(
+      cmlp, x, x2, console_kwargs=CONSOLE_TEST_KWARGS
+    ).splitlines()
+
+    self.assertIn('CustomMLP Summary', table_repr[0])
+    self.assertIn('float32[1,4]', table_repr[4])
+    self.assertIn('float32[1,8]', table_repr[5])
+    self.assertIn('float32[1,8]', table_repr[6])
+
+
+if __name__ == '__main__':
+  absltest.main()
\ No newline at end of file