improves float check + adds tests

Author: JyotinderSingh

keras_hub/src/models/task_test.py

@@ -163,7 +163,7 @@ def test_quantized_preset_loading_and_saving(
                     layer.kernel.dtype,
                     expected_dtype,
                     f"Layer '{layer.name}' kernel "
-                    "should have dtype '{expected_dtype}'",
+                    f"should have dtype '{expected_dtype}'",
                 )
 
         # Ensure inference runs without errors.

keras_hub/src/utils/preset_utils.py

@@ -723,11 +723,12 @@ def _resolve_dtype(self, config, kwargs):
           saved types verbatim.
 
         Args:
-            config: The model configuration.
-            kwargs: Additional keyword arguments, potentially including `dtype`.
+            config: dict. The model configuration.
+            kwargs: dict. Additional keyword arguments, potentially including
+             `dtype`.
 
         Returns:
-            The resolved dtype.
+            str, dict, or DTypePolicy. The resolved dtype.
         """
         # 1. If a user specified dtype is passed, use that.
         if "dtype" in kwargs and kwargs["dtype"] is not None:

keras_hub/src/utils/tensor_utils.py

@@ -310,7 +310,29 @@ def is_tensor_type(x):
 
 
 def is_float_dtype(dtype):
-    return "float" in keras.backend.standardize_dtype(dtype)
+    """
+    Checks if a dtype is a float type by using a regex.
+
+    This function standardizes the input dtype and then uses a regular
+    expression to perform an exact match. It identifies standard floats,
+    bfloats, and mixed-precision float types.
+
+    For example:
+    - `is_float_dtype("float32")` returns `True`.
+    - `is_float_dtype("bfloat16")` returns `True`.
+    - `is_float_dtype("mixed_float16")` returns `True`.
+    - `is_float_dtype("int8")` returns `False`.
+    - `is_float_dtype("int8_from_float32")` returns `False`.
+
+    Args:
+        dtype: str, DTypePolicy. The data type to check.
+
+    Returns:
+        bool: `True` if the dtype is a floating-point type, `False` otherwise.
+    """
+    pattern = re.compile(r"^(mixed_)?(b)?float[0-9]*$")
+    standardized_dtype = keras.backend.standardize_dtype(dtype)
+    return pattern.match(standardized_dtype) is not None
 
 
 def is_int_dtype(dtype):

keras_hub/src/utils/tensor_utils_test.py

@@ -8,6 +8,7 @@
 from keras_hub.src.utils.tensor_utils import convert_preprocessing_inputs
 from keras_hub.src.utils.tensor_utils import convert_preprocessing_outputs
 from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
+from keras_hub.src.utils.tensor_utils import is_float_dtype
 from keras_hub.src.utils.tensor_utils import is_tensor_type
 from keras_hub.src.utils.tensor_utils import preprocessing_function
 from keras_hub.src.utils.tensor_utils import target_gather
@@ -304,3 +305,28 @@ def test_target_gather_invalid_rank(self):
         indices = np.array([0, 1], dtype="int32")
         with self.assertRaisesRegex(ValueError, "larger than 3"):
             _ = target_gather(targets, indices)
+
+
+class IsFloatDtypeTest(TestCase):
+    def test_float_dtypes_return_true(self):
+        float_dtypes = [
+            "float16",
+            "float32",
+            "float64",
+            "bfloat16",
+        ]
+        for dtype in float_dtypes:
+            self.assertTrue(is_float_dtype(dtype))
+
+    def test_non_float_dtypes_return_false(self):
+        non_float_dtypes = [
+            "int8",
+            "int32",
+            "uint8",
+            "bool",
+            "string",
+            "int8_from_float32",
+            "int4_from_bfloat16",
+        ]
+        for dtype in non_float_dtypes:
+            self.assertFalse(is_float_dtype(dtype))