Merge pull request #3452 from chiamp:tabulate

PiperOrigin-RevId: 580696478
google · Nov 9, 2023 · 50de4c4 · 50de4c4
2 parents b6eca17 + e60eda3
commit 50de4c4
Showing 1 changed file with 24 additions and 24 deletions.
diff --git a/flax/linen/module.py b/flax/linen/module.py
@@ -2434,30 +2434,30 @@ def __call__(self, x):
 
     This gives the following output::
 
-                                           Foo Summary
-    ┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓
-    ┃ path    ┃ module ┃ inputs        ┃ outputs       ┃ flops ┃ vjp_flops ┃ params          ┃
-    ┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩
-    │         │ Foo    │ float32[16,9] │ float32[16,2] │ 1504  │ 4460      │                 │
-    ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤
-    │ Dense_0 │ Dense  │ float32[16,9] │ float32[16,4] │ 1216  │ 3620      │ bias:           │
-    │         │        │               │               │       │           │ float32[4]      │
-    │         │        │               │               │       │           │ kernel:         │
-    │         │        │               │               │       │           │ float32[9,4]    │
-    │         │        │               │               │       │           │                 │
-    │         │        │               │               │       │           │ 40 (160 B)      │
-    ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤
-    │ Dense_1 │ Dense  │ float32[16,4] │ float32[16,2] │ 288   │ 840       │ bias:           │
-    │         │        │               │               │       │           │ float32[2]      │
-    │         │        │               │               │       │           │ kernel:         │
-    │         │        │               │               │       │           │ float32[4,2]    │
-    │         │        │               │               │       │           │                 │
-    │         │        │               │               │       │           │ 10 (40 B)       │
-    ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤
-    │         │        │               │               │       │     Total │ 50 (200 B)      │
-    └─────────┴────────┴───────────────┴───────────────┴───────┴───────────┴─────────────────┘
-
-                                   Total Parameters: 50 (200 B)
+                                            Foo Summary
+      ┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓
+      ┃ path    ┃ module ┃ inputs        ┃ outputs       ┃ flops ┃ vjp_flops ┃ params          ┃
+      ┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩
+      │         │ Foo    │ float32[16,9] │ float32[16,2] │ 1504  │ 4460      │                 │
+      ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤
+      │ Dense_0 │ Dense  │ float32[16,9] │ float32[16,4] │ 1216  │ 3620      │ bias:           │
+      │         │        │               │               │       │           │ float32[4]      │
+      │         │        │               │               │       │           │ kernel:         │
+      │         │        │               │               │       │           │ float32[9,4]    │
+      │         │        │               │               │       │           │                 │
+      │         │        │               │               │       │           │ 40 (160 B)      │
+      ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤
+      │ Dense_1 │ Dense  │ float32[16,4] │ float32[16,2] │ 288   │ 840       │ bias:           │
+      │         │        │               │               │       │           │ float32[2]      │
+      │         │        │               │               │       │           │ kernel:         │
+      │         │        │               │               │       │           │ float32[4,2]    │
+      │         │        │               │               │       │           │                 │
+      │         │        │               │               │       │           │ 10 (40 B)       │
+      ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤
+      │         │        │               │               │       │     Total │ 50 (200 B)      │
+      └─────────┴────────┴───────────────┴───────────────┴───────┴───────────┴─────────────────┘
+
+                                    Total Parameters: 50 (200 B)
 
     **Note**: rows order in the table does not represent execution order,
     instead it aligns with the order of keys in `variables` which are sorted