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:

Fix (patch):
diff --git a/flax/nnx/summary.py b/flax/nnx/summary.py
index 7f1aed53..7396b2ed 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:


Test changes:
diff --git a/tests/nnx/summary_test.py b/tests/nnx/summary_test.py
index 94059c21..2555fde2 100644
--- a/tests/nnx/summary_test.py
+++ b/tests/nnx/summary_test.py
@@ -72,3 +72,31 @@ class SummaryTest(absltest.TestCase):
     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