Fix tf dataset detection logic. (#21794)

In some contexts, matching the begining of the module name won't work, we need to match the tensorflow module in any part of the module name.

This was already done in other contexts:

https://github.com/keras-team/keras/blob/0512fdb3c70c6ef5499615793e40d5ad46f3b301/keras/src/trainers/data_adapters/data_adapter_utils.py#L307

Author: hertschuh

keras/src/utils/dataset_utils.py

@@ -461,7 +461,7 @@ def is_tf_dataset(dataset):
     return _mro_matches(
         dataset,
         class_names=("DatasetV2", "Dataset"),
-        module_prefixes=(
+        module_substrings=(
             "tensorflow.python.data",  # TF classic
             "tensorflow.data",  # newer TF paths
         ),
@@ -480,14 +480,18 @@ def is_torch_dataset(dataset):
     return _mro_matches(dataset, ("Dataset",), ("torch.utils.data",))
 
 
-def _mro_matches(dataset, class_names, module_prefixes):
+def _mro_matches(
+    dataset, class_names, module_prefixes=(), module_substrings=()
+):
     if not hasattr(dataset, "__class__"):
         return False
     for parent in dataset.__class__.__mro__:
         if parent.__name__ in class_names:
             mod = str(parent.__module__)
             if any(mod.startswith(pref) for pref in module_prefixes):
                 return True
+            if any(subs in mod for subs in module_substrings):
+                return True
     return False