Refactor variable serialization.

[james77777778]

This is an alternative solution compared to #21681

Since the default behavior is to store the variables using their order as the keys, we can maintain this behavior by using a variable_serialization_spec to explicitly define the order.
variable_serialization_spec works like a human-readable proxy between the variables and store for better compatibility.

IMO, save_own_variables and load_own_variables become more readable and simpler in this PR.

cc @hertschuh @fchollet @JyotinderSingh

[gemini-code-assist]

Summary of Changes

Hello @james77777778, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request implements a significant refactoring of how Keras layers handle variable serialization, especially concerning quantized models. By introducing a variable_serialization_spec, the PR establishes a clear, explicit, and ordered mechanism for saving and loading layer variables. This change simplifies the internal logic of save_own_variables and load_own_variables methods in layers like Dense, EinsumDense, and Embedding, making the process more robust and maintainable. It also centralizes common variable checking logic in the base Layer class, improving code organization and reducing redundancy.

Highlights

• Unified Variable Serialization: Introduced variable_serialization_spec to explicitly define the ordered list of variables for serialization across different quantization modes in Dense, EinsumDense, and Embedding layers.
• Simplified Save/Load Logic: The save_own_variables and load_own_variables methods in affected layers are refactored to iterate through the new variable_serialization_spec, using numerical indices as keys for storing and retrieving variables, leading to more readable and simpler implementations.
• Centralized Variable Check: The _check_load_own_variables utility method, which validates the number of variables during loading, has been moved from individual layer implementations to the base keras.src.layers.layer.Layer class, promoting code reuse and consistency.
• Legacy Code Removal: The _legacy_load_own_variables method, previously used for handling older serialization formats, has been removed from Dense and EinsumDense layers, streamlining the loading process.
• Test Updates: Test cases in file_editor_test.py were updated to reflect the new numerical indexing scheme for variable weights, ensuring compatibility with the refactored serialization.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a solid refactoring of variable serialization in Dense, EinsumDense, and Embedding layers. By introducing variable_serialization_spec, the serialization logic becomes more explicit, readable, and consistent, which is a great improvement for maintainability. The removal of legacy loading paths and the consolidation of checks into the base Layer class further streamline the code. The changes are well-implemented and align with the goal of simplifying the serialization process. I have one minor suggestion regarding a comment to improve clarity.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 87.30159% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.63%. Comparing base (3fac66f) to head (9615b5b).

Files with missing lines	Patch %	Lines
keras/src/layers/core/embedding.py	76.47%	1 Missing and 3 partials ⚠️
keras/src/layers/core/dense.py	90.47%	0 Missing and 2 partials ⚠️
keras/src/layers/core/einsum_dense.py	90.47%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21713      +/-   ##
==========================================
+ Coverage   82.60%   82.63%   +0.03%     
==========================================
  Files         572      572              
  Lines       58326    58284      -42     
  Branches     9134     9117      -17     
==========================================
- Hits        48179    48164      -15     
+ Misses       7817     7802      -15     
+ Partials     2330     2318      -12     

Flag	Coverage Δ
keras	82.43% <87.30%> (+0.03%)	⬆️
keras-jax	63.33% <87.30%> (+0.01%)	⬆️
keras-numpy	57.67% <74.60%> (+0.01%)	⬆️
keras-openvino	34.33% <11.11%> (+0.01%)	⬆️
keras-tensorflow	64.07% <87.30%> (+0.02%)	⬆️
keras-torch	63.66% <87.30%> (+0.02%)	⬆️

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

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[fchollet]

Thanks for the PR!

To be clear, the reason we were using indices instead of string names was that string names are susceptible to changes, while indices are more stable. It makes the format more future proof.

One thing I've considered doing (and which could be done across the board, for all weights in all layers) is to have names like "0:kernel", "1:bias", which would be parsed as int(name.split(":")[0]). This would be human readable and fully backwards compatible with the old format. What do you think?

[hertschuh]

Thanks for the PR!

To be clear, the reason we were using indices instead of string names was that string names are susceptible to changes, while indices are more stable. It makes the format more future proof.

Thanks for providing the context.

One thing I've considered doing (and which could be done across the board, for all weights in all layers) is to have names like "0:kernel", "1:bias", which would be parsed as int(name.split(":")[0]). This would be human readable and fully backwards compatible with the old format. What do you think?

Unfortunately, this is not backwards compatible in the way we were discussing.

Let's say we save the latest Gemma checkpoints with Keras 3.12 with this new format, people using Keras 3.11 won't be able to load them.

[fchollet]

Let's say we save the latest Gemma checkpoints with Keras 3.12 with this new format, people using Keras 3.11 won't be able to load them.

Yes, but historically it has been the case that file loading has been backwards compatible (newer Keras versions can load older files) but not forward compatible (older Keras versions may not necessarily load newer files).

The only alternative here is to keep integer indices.