diff --git a/flax/linen/module.py b/flax/linen/module.py index 37175e7df8..926cfc7825 100644 --- 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