Add optimizer weights to the model summary.

[qlzh727]

The sample output with optimizer weights is like below:

Model: "language_model"
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)        ┃ Output Shape      ┃ Param # ┃ Connected to         ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ padding_mask        │ (None, None)      │       0 │ -                    │
│ (InputLayer)        │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ token_ids           │ (None, None)      │       0 │ -                    │
│ (InputLayer)        │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ opt_backbone        │ (None, None,      │ 1,315,… │ padding_mask[0][0],  │
│ (OPTBackbone)       │ 2048)             │         │ token_ids[0][0]      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ token_embedding     │ (None, None,      │ 102,95… │ opt_backbone[0][0]   │
│ (ReversibleEmbeddi… │ 50272)            │         │                      │
└─────────────────────┴───────────────────┴─────────┴──────────────────────┘
 Total params: 3,947,261,954 (14.70 GB)
 Trainable params: 1,315,753,984 (4.90 GB)
 Non-trainable params: 0 (0.00 B)
 Optimizer params: 2,631,507,970 (9.80 GB)

[codecov]

Codecov Report

Patch coverage: 100.00% and project coverage change: +0.14% 🎉

Comparison is base (d0b53fd) 76.10% compared to head (1689287) 76.25%.
Report is 4 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #859      +/-   ##
==========================================
+ Coverage   76.10%   76.25%   +0.14%     
==========================================
  Files         329      329              
  Lines       31270    31276       +6     
  Branches     6083     6084       +1     
==========================================
+ Hits        23799    23850      +51     
+ Misses       5871     5834      -37     
+ Partials     1600     1592       -8     

Flag	Coverage Δ
keras_core	76.16% <100.00%> (+0.14%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files Changed	Coverage Δ
keras_core/utils/summary_utils.py	72.72% <100.00%> (+5.23%)	⬆️

... and 4 files with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[fchollet]

Thanks for the PR -- looks good! Please add a unit test.

[qlzh727]

Thanks for the PR -- looks good! Please add a unit test.

Sure. Done.

[fchollet]

LGTM, thanks!