switch to re.fullmatch for more explicit behavior spec

Author: JyotinderSingh

keras/src/dtype_policies/dtype_policy_map.py

@@ -98,25 +98,22 @@ def quantization_mode(self):
         return self.default_policy.quantization_mode
 
     def __getitem__(self, key):
-        """Retrieves a `DTypePolicy` by its key, with a regex fallback.
+        """Retrieves a `DTypePolicy` by treating map keys as regex patterns.
 
-        This method first attempts an exact key match.
-        If no exact match is found, it treats the keys in the map as regular
-        expression patterns. A pattern is considered a match only if it aligns
-        with the start of the input `key` and covers a complete path component
-        (i.e., it is followed by a `/` or the end of the string).
+        This method first attempts an exact key match for performance. 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, a map key "encoder" will match the path
-        "encoder/attention", but it will not match "encoder_v2". This
-        component-wise matching prevents incorrect policy assignments to
-        similarly named layers.
+        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: The string key to query for a `DTypePolicy`.
+            key: str. The key to query for a `DTypePolicy`.
 
         Returns:
-            The corresponding `DTypePolicy`. If no valid match is found
-            (either exact or regex), this method returns `self.default_policy`.
+            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
@@ -126,14 +123,11 @@ def __getitem__(self, key):
         if key in self._policy_map:
             return self._policy_map[key]
 
-        # 2. If no exact match is found, fallback to a regex match.
-        # Check for a match that covers a full path component.
-        # The pattern must match from the start of the `key` and be
-        # followed by either a '/' or the end of the string.
+        # 2. Fallback to a full regex match.
         matching_keys = [
             pattern
             for pattern in self._policy_map
-            if re.match(f"{pattern}(/|$)", key)
+            if re.fullmatch(pattern, key)
         ]
 
         # 3. Handle cases based on the number of matches found.

keras/src/dtype_policies/dtype_policy_map_test.py

@@ -135,49 +135,47 @@ def test_get(self):
         policy_map = DTypePolicyMap()
         # Policy for an exact layer path
         policy_map["model/encoder/layer_0/dense"] = bfloat16_policy
-        # Policy that could be partially matched
+        # 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
-        policy_map["model/decoder"] = float32_policy
-        policy_map["model/decoder/attention"] = float16_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
-        # An exact key should always return the correct policy.
         self.assertEqual(
             policy_map["model/encoder/layer_0/dense"], bfloat16_policy
         )
         self.assertEqual(
             policy_map["model/encoder/attention/query"], int8_policy
         )
 
-        # 3. Test successful regex fallback (component-wise match)
-        # "model/decoder" should match "model/decoder/layer_0" because
-        # it's a full component prefix.
+        # 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 prevention of partial regex match
-        # "model/encoder/attention/query" should NOT match
-        # "model/encoder/attention/query_norm"
-        # as it's not followed by a '/' or the end of the string.
+        # 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)
 
         # 5. Test no match
-        # A key with no exact or valid regex match should return the default.
         self.assertEqual(
             policy_map["model/embedding"], policy_map.default_policy
         )
-        self.assertEqual(
-            policy_map["prefix/model/decoder/attention"],
-            policy_map.default_policy,
-        )
 
         # 6. Test multiple regex matches causing a ValueError
-        # The path "model/decoder/attention/output" matches two regex keys:
-        # - "model/decoder"
-        # - "model/decoder/attention"
+        # "model/decoder/attention/output" matches two regex patterns:
+        # - "model/decoder/.*"
+        # - "model/decoder/attention/.*"
         with self.assertRaisesRegex(
             ValueError,
             "Path 'model/decoder/attention/output' matches multiple "