Ignore regex fallback for quantized policies in DTypePolicyMap (#21608)

* Ignore regex match for quantized policies in DTypePolicyMap

* only match complete path components when falling back to regex match

* updates docstring

* minor conditional modification

* updates comment

* replace re.search with re.match

* switch to re.fullmatch for more explicit behavior spec

* Added detailed example to docstring

Author: JyotinderSingh

keras/src/dtype_policies/dtype_policy_map.py

@@ -35,29 +35,54 @@ def get_config(self):
     However, it is also possible to set a regex as the key. See the docstring of
     `get` for more details.
 
-    See below for a usage example. You can define the naming schema
-    of the `DTypePolicy`, and then retrieve the corresponding `DTypePolicy`
-    instance.
-
-    ```python
-    dtype_policy_map = DTypePolicyMap()
-    dtype_policy_map["layer/dense_0"] = DTypePolicy("bfloat16")
-    dtype_policy_map["layer/dense_1"] = QuantizedDTypePolicy("int8", "bfloat16")
-
-    policy_0 = dtype_policy_map["layer/dense_0"]
-    policy_1 = dtype_policy_map["layer/dense_1"]
-    policy_2 = dtype_policy_map["layer/dense_2"]  # No hit
-    assert policy_0 == DTypePolicy("bfloat16")
-    assert policy_1 == QuantizedDTypePolicy("int8", "bfloat16")
-    assert policy_2 == keras.config.dtype_policy()
-    ```
-
     Args:
         default_policy: An optional `DTypePolicy` instance specifying the
             default dtype policy. If not specified, the value will default to
             `keras.config.dtype_policy()`.
         policy_map: An optional dict that maps string to `DTypePolicy`
             instances. Defaults to `None`
+
+    Example:
+
+    ```python
+    >>> from keras.src import dtype_policies
+    >>> bfloat16 = dtype_policies.DTypePolicy("bfloat16")
+    >>> float16 = dtype_policies.DTypePolicy("float16")
+    >>> float32 = dtype_policies.DTypePolicy("float32")
+    >>> policy_map = DTypePolicyMap(default_policy=float32)
+
+    # Set policies using an exact path and a regex pattern.
+    # Note: "decoder" will only match the exact path, not its children.
+    >>> policy_map["encoder/layer_0/dense"] = bfloat16
+    >>> policy_map["encoder/.*"] = float16
+    >>> policy_map["decoder"] = bfloat16
+
+    # 1. An exact match is found and returned directly.
+    >>> policy_map["encoder/layer_0/dense"].name
+    'bfloat16'
+
+    # 2. A regex match is found for a child layer.
+    # It matches the "encoder/.*" pattern.
+    >>> policy_map["encoder/attention/query"].name
+    'float16'
+
+    # 3. No implicit prefix matching occurs.
+    # "decoder/attention" does not match the key "decoder".
+    # The default policy is returned.
+    >>> policy_map["decoder/attention"].name
+    'float32'
+
+    # 4. A ValueError is raised if a path matches multiple patterns.
+    >>> policy_map["encoder/attention/.*"] = bfloat16
+    # "encoder/attention/query" now matches two patterns:
+    # - "encoder/.*"
+    # - "encoder/attention/.*"
+    >>> try:
+    ...     policy_map["encoder/attention/query"]
+    ... except ValueError as e:
+    ...     print(e)
+    Path 'encoder/attention/query' matches multiple dtype policy ..
+    ```
     """
 
     def __init__(self, default_policy=None, policy_map=None):
@@ -100,24 +125,79 @@ def quantization_mode(self):
     def __getitem__(self, key):
         """Retrieves the corresponding `DTypePolicy` by the string key.
 
-        When there isn't an exact match, all the existing keys in the map
-        will be treated as a regex and map against the input key again. When
-        there are multiple matches for the regex, an `ValueError` will be
-        raised. Returns `self.default_policy` if there isn't any match found.
+        This method first attempts an exact key match. If no exact match is
+        found, it treats all keys in the map as regular expression patterns
+        and uses `re.fullmatch` to find a policy.
+
+        For example, to apply a policy to all sublayers of an `encoder` block,
+        the key should be explicitly set to `"encoder/.*"`. A key of
+        `"encoder"` will only match the layer with that exact path.
 
         Args:
-            key: String key to query a `DTypePolicy`.
+            key: str. The key to query for a `DTypePolicy`.
 
         Returns:
-            Corresponding `DTypePolicy` based on the query.
+            The corresponding `DTypePolicy`. If no match is found, this method
+            returns `self.default_policy`.
+
+        Raises:
+            ValueError: If the `key` matches more than one regex pattern in the
+            map.
+
+        Example:
+
+        ```python
+        >>> from keras.src import dtype_policies
+        >>> bfloat16 = dtype_policies.DTypePolicy("bfloat16")
+        >>> float16 = dtype_policies.DTypePolicy("float16")
+        >>> float32 = dtype_policies.DTypePolicy("float32")
+        >>> policy_map = DTypePolicyMap(default_policy=float32)
+
+        # Set policies using an exact path and a regex pattern.
+        # Note: "decoder" will only match the exact path, not its children.
+        >>> policy_map["encoder/layer_0/dense"] = bfloat16
+        >>> policy_map["encoder/.*"] = float16
+        >>> policy_map["decoder"] = bfloat16
+
+        # 1. An exact match is found and returned directly.
+        >>> policy_map["encoder/layer_0/dense"].name
+        'bfloat16'
+
+        # 2. A regex match is found for a child layer.
+        # It matches the "encoder/.*" pattern.
+        >>> policy_map["encoder/attention/query"].name
+        'float16'
+
+        # 3. No implicit prefix matching occurs.
+        # "decoder/attention" does not match the key "decoder".
+        # The default policy is returned.
+        >>> policy_map["decoder/attention"].name
+        'float32'
+
+        # 4. A ValueError is raised if a path matches multiple patterns.
+        >>> policy_map["encoder/attention/.*"] = bfloat16
+        # "encoder/attention/query" now matches two patterns:
+        # - "encoder/.*"
+        # - "encoder/attention/.*"
+        >>> try:
+        ...     policy_map["encoder/attention/query"]
+        ... except ValueError as e:
+        ...     print(e)
+        Path 'encoder/attention/query' matches multiple dtype policy ..
+        ```
         """
+        # 1. Check for an exact match.
         if key in self._policy_map:
             return self._policy_map[key]
 
-        matching_keys = []
-        for k in self._policy_map:
-            if re.search(k, key):
-                matching_keys.append(k)
+        # 2. Fallback to a full regex match.
+        matching_keys = [
+            pattern
+            for pattern in self._policy_map
+            if re.fullmatch(pattern, key)
+        ]
+
+        # 3. Handle cases based on the number of matches found.
         if len(matching_keys) > 1:
             raise ValueError(
                 f"Path '{key}' matches multiple dtype policy "
@@ -127,6 +207,8 @@ def __getitem__(self, key):
             )
         elif len(matching_keys) == 1:
             return self._policy_map[matching_keys[0]]
+
+        # 4. If there were no matches, return the default.
         return self.default_policy
 
     def __setitem__(self, key, policy):

keras/src/dtype_policies/dtype_policy_map_test.py

@@ -124,50 +124,63 @@ def test_add(self):
             dtype_policy_map["layer/dense_3"] = 123
 
     def test_get(self):
-        dtype_policy_map = DTypePolicyMap()
-        dtype_policy_map["layer/dense_0"] = dtype_policies.DTypePolicy(
-            "bfloat16"
-        )
-        dtype_policy_map["layer/dense_1"] = dtype_policies.QuantizedDTypePolicy(
+        # 1. Setup
+        bfloat16_policy = dtype_policies.DTypePolicy("bfloat16")
+        int8_policy = dtype_policies.QuantizedDTypePolicy(
             "int8", "mixed_bfloat16"
         )
-        dtype_policy_map["layer/dense_2"] = (
-            dtype_policies.QuantizedFloat8DTypePolicy("float8", "mixed_float16")
-        )
+        float32_policy = dtype_policies.DTypePolicy("float32")
+        float16_policy = dtype_policies.DTypePolicy("float16")
 
+        policy_map = DTypePolicyMap()
+        # Policy for an exact layer path
+        policy_map["model/encoder/layer_0/dense"] = bfloat16_policy
+        # Policy for a layer that is also a prefix of another layer's name
+        policy_map["model/encoder/attention/query"] = int8_policy
+        # Regex policies for entire scopes MUST include wildcards
+        policy_map["model/decoder/.*"] = float32_policy
+        policy_map["model/decoder/attention/.*"] = float16_policy
+
+        # 2. Test exact match
         self.assertEqual(
-            dtype_policy_map["layer/dense_0"],
-            dtype_policies.DTypePolicy("bfloat16"),
-        )
-        self.assertEqual(
-            dtype_policy_map["layer/dense_1"],
-            dtype_policies.QuantizedDTypePolicy("int8", "mixed_bfloat16"),
+            policy_map["model/encoder/layer_0/dense"], bfloat16_policy
         )
         self.assertEqual(
-            dtype_policy_map["layer/dense_2"],
-            dtype_policies.QuantizedFloat8DTypePolicy(
-                "float8", "mixed_float16"
-            ),
+            policy_map["model/encoder/attention/query"], int8_policy
         )
 
-        self.assertNotEqual(
-            dtype_policy_map["layer/dense_2"],
-            dtype_policies.QuantizedFloat8DTypePolicy("float8", "bfloat16"),
+        # 3. Test successful regex fallback (explicit wildcard)
+        # "model/decoder/.*" should match its children.
+        self.assertEqual(policy_map["model/decoder/layer_0"], float32_policy)
+
+        # 4. Test that partial matches are ignored
+        # The exact key "model/encoder/attention/query" should not match
+        # "model/encoder/attention/query_norm" without a wildcard.
+        self.assertEqual(
+            policy_map["model/encoder/attention/query_norm"],
+            policy_map.default_policy,
         )
+        # A plain key "model/decoder" will not match "model/decoder/layer_0"
+        policy_map["model/decoder"] = bfloat16_policy  # Add exact key
+        self.assertEqual(policy_map["model/decoder/layer_0"], float32_policy)
+        # Still matches the more general regex
+        self.assertEqual(policy_map["model/decoder"], bfloat16_policy)
 
-        # No hit
+        # 5. Test no match
         self.assertEqual(
-            dtype_policy_map["layer/batch_normalization"],
-            dtype_policy_map.default_policy,
+            policy_map["model/embedding"], policy_map.default_policy
         )
 
-        # It will cause a ValueError in the case of one-to-many.
-        dtype_policy_map["dense"] = dtype_policies.DTypePolicy("float32")
-        dtype_policy_map["dense_1"] = dtype_policies.DTypePolicy("float32")
+        # 6. Test multiple regex matches causing a ValueError
+        # "model/decoder/attention/output" matches two regex patterns:
+        # - "model/decoder/.*"
+        # - "model/decoder/attention/.*"
         with self.assertRaisesRegex(
-            ValueError, "Path 'dense_10' matches multiple dtype policy"
+            ValueError,
+            "Path 'model/decoder/attention/output' matches multiple "
+            "dtype policy",
         ):
-            dtype_policy_map["dense_10"]
+            _ = policy_map["model/decoder/attention/output"]
 
     def test_delete(self):
         dtype_policy_map = DTypePolicyMap()