Merge pull request #4580 from google:nnx-tabulate-issue

PiperOrigin-RevId: 731490914

Author: Flax Authors

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:

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()