Ignore regex fallback for quantized policies in DTypePolicyMap

[JyotinderSingh]

Issue

The regex fallback logic in DTypePolicyMap.__getitem__ could incorrectly assign a quantized policy to a non-quantized layer.

This issue arises when deserializing models, particularly from KerasHub, where policies for quantized layers are stored in a DTypePolicyMap. The keys for these policies are intended for exact matches. However, the regex fallback can cause a non-quantized layer's path (e.g., "attention/query_norm") to erroneously match a key associated with a quantized policy (e.g., "attention/query").

Consequently, the non-quantized layer receives the wrong policy and fails during the forward pass when it attempts to call a quantized_call method it doesn't have.

Fix

Changed the fallback logic from implicit component matching to explicit regex matching using re.fullmatch. This makes the matching behavior more predictable. Keys in the policy map are now treated as full regular expression patterns that must match the entire layer path.

This means:

1. A key of "attention/query" will only match a layer with the exact path "attention/query".
2. It will not match related paths like "attention/query/sub_layer" or "attention/query_norm".
3. To match all sublayers within a scope, an explicit wildcard must be used and stored in the map, e.g., "attention/.*".

Testing

Modified the existing tests to cover all possible cases.

Colab

1. Failing Case: This notebook shows that a quantized gemma3_1b KerasHub model fails to perform inference after being saved and loaded. The happens since the the query_norm (an RMSNormalization layer) matches the regex for another quantized layer called query (an EinsumDense layer for calculating query values) when resolving the dtype.
2. Successful Case: This notebook shows that the above issue no longer occurs after this fix is applied.

[gemini-code-assist]

Summary of Changes

Hello @JyotinderSingh, 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 addresses a bug in DTypePolicyMap where its regex fallback mechanism could erroneously assign a quantized policy to a non-quantized layer, particularly during model deserialization from sources like KerasHub. The fix modifies the getitem method to ensure that if a policy is found via regex matching, it is only applied if it is not a quantized policy; otherwise, the default policy is used. This prevents incorrect policy assignments and subsequent runtime failures.

Highlights

• Bug Fix: Resolved an issue in DTypePolicyMap.getitem where regex fallback could incorrectly assign quantized policies to non-quantized layers.
• Policy Assignment Logic: Modified the getitem method to explicitly check if a policy obtained via regex matching is quantized. If it is, the default policy is returned instead.
• Impact: Prevents runtime failures in non-quantized layers that might otherwise receive an incompatible quantized policy, especially when deserializing models.
• Testing: Added two new unit tests (test_get_with_regex_non_quantized and test_get_with_regex_quantized) to dtype_policy_map_test.py to validate the updated policy retrieval logic.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 provides a well-reasoned and cleanly implemented fix for an issue where the regex fallback in DTypePolicyMap could incorrectly assign a quantized policy to a non-quantized layer. The logic added to __getitem__ to check if a regex-matched policy is quantized is a direct and effective solution. The code is made more maintainable by the addition of clear comments and an updated docstring that explains the behavior. The new unit tests in dtype_policy_map_test.py are excellent, as they specifically target and validate both the intended new behavior and the case it prevents. The changes are of high quality and improve the robustness of the system.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.45%. Comparing base (19367bc) to head (e8e528d).
⚠️ Report is 3 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21608      +/-   ##
==========================================
- Coverage   82.73%   82.45%   -0.29%     
==========================================
  Files         572      572              
  Lines       57261    57337      +76     
  Branches     8961     8969       +8     
==========================================
- Hits        47377    47275     -102     
- Misses       7672     7762      +90     
- Partials     2212     2300      +88     

Flag	Coverage Δ
keras	82.25% <100.00%> (-0.29%)	⬇️
keras-jax	63.57% <100.00%> (-0.21%)	⬇️
keras-numpy	57.86% <100.00%> (-0.15%)	⬇️
keras-openvino	34.34% <0.00%> (-0.08%)	⬇️
keras-tensorflow	64.21% <100.00%> (-0.15%)	⬇️
keras-torch	63.78% <100.00%> (-0.14%)	⬇️

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.

[mattdangerw]

awesome! thank you