PyPI - keras-hub-nightly - Versions diffs - 0.20.0.dev202504020401__tar.gz → 0.20.0.dev202504030357__tar.gz - Mend

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-hub-nightly
-Version: 0.20.0.dev202504020401
+Version: 0.20.0.dev202504030357
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team

@@ -6,7 +6,9 @@ from keras import ops
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
+from keras_hub.src.utils.keras_utils import gpu_supports_fused_attention_op
+from keras_hub.src.utils.keras_utils import running_on_gpu
 from keras_hub.src.utils.keras_utils import running_on_tpu
@@ -106,17 +108,22 @@ class CachedGemmaAttention(keras.layers.Layer):
         )
         return x
-    def _can_use_flash_attention(self):
-        if not has_flash_attention_support():
+    def _use_fused_attention_op(self):
+        if not fused_attention_op_available():
             return False
         if self.dropout > 0.0:
             return False
-        if self.logit_soft_cap is None:
-            return True
-        sig = inspect.signature(ops.dot_product_attention)
-        # We can currently only run soft capped attention for keras >= 3.10
-        # and only on TPU.
-        return running_on_tpu() and "attn_logits_soft_cap" in sig.parameters
+        if running_on_gpu():
+            # GPU never supports softcap in the fused op.
+            if self.logit_soft_cap is not None:
+                return False
+            return gpu_supports_fused_attention_op()
+        elif running_on_tpu():
+            # TPU supports softcap with on keras >= 3.10.
+            sig = inspect.signature(ops.dot_product_attention)
+            return "attn_logits_soft_cap" in sig.parameters
+        else:
+            return False
     def _compute_attention(
         self,
@@ -140,7 +147,7 @@ class CachedGemmaAttention(keras.layers.Layer):
                 cache_update_index=cache_update_index,
             )
-        if self._can_use_flash_attention():
+        if self._use_fused_attention_op():
             if attention_mask is not None:
                 attention_mask = ops.expand_dims(attention_mask, axis=1)
                 attention_mask = ops.cast(attention_mask, dtype="bool")

@@ -7,7 +7,7 @@ from keras import ops
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.models.gemma.rms_normalization import RMSNormalization
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 from keras_hub.src.utils.keras_utils import running_on_tpu
@@ -140,7 +140,7 @@ class CachedGemma3Attention(keras.layers.Layer):
         return x
     def _can_use_flash_attention(self):
-        if not has_flash_attention_support():
+        if not fused_attention_op_available():
             return False
         if self.dropout > 0.0:
             return False

@@ -5,7 +5,7 @@ from keras import ops
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 class GPTNeoXAttention(keras.layers.Layer):
@@ -125,7 +125,7 @@ class GPTNeoXAttention(keras.layers.Layer):
     def _compute_attention(
         self, query, key, value, attention_mask=None, training=None
     ):
-        if has_flash_attention_support() and self.dropout == 0:
+        if fused_attention_op_available() and self.dropout == 0:
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

@@ -5,7 +5,7 @@ from keras import ops
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 class LlamaAttention(keras.layers.Layer):
@@ -185,7 +185,7 @@ class LlamaAttention(keras.layers.Layer):
         return self._softmax(attention_scores)
     def _compute_attention(self, query, key, value, attention_mask=None):
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

@@ -6,8 +6,8 @@ from keras import ops
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
 from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 from keras_hub.src.utils.keras_utils import gelu_approximate
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
 from keras_hub.src.utils.keras_utils import standardize_data_format
@@ -771,7 +771,7 @@ class MMDiTBlock(layers.Layer):
     def _compute_attention(self, query, key, value):
         batch_size = ops.shape(query)[0]
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             encoded = ops.dot_product_attention(

@@ -55,7 +55,7 @@ def standardize_data_format(data_format):
     return data_format
-def has_flash_attention_support():
+def fused_attention_op_available():
     if (
         hasattr(keras.config, "is_flash_attention_enabled")
         and keras.config.backend() == "jax"
@@ -104,3 +104,46 @@ def running_on_gpu():
         import torch
         return torch.cuda.is_available()
+def gpu_supports_fused_attention_op():
+    deny_list = ["T4"]
+    for denied_gpu in deny_list:
+        if any(denied_gpu in gpu.upper() for gpu in get_gpu_names()):
+            return False
+    return True
+def get_gpu_names():
+    """Detects and returns the names of available GPUs based on the backend.
+    Note:
+        The format and content of the returned GPU names are **not normalized**
+        and vary significantly depending on the active backend. This function
+        provides the names as reported by the respective backend's API."
+    """
+    backend = keras.config.backend()
+    if backend == "jax":
+        import jax
+        devices = jax.devices()
+        return [getattr(d, "device_kind", "") for d in devices]
+    elif backend == "tensorflow":
+        import tensorflow as tf
+        gpus = tf.config.list_physical_devices("GPU")
+        return [
+            tf.config.experimental.get_device_details(gpu)["device_name"]
+            for gpu in gpus
+        ]
+    elif backend == "torch":
+        import torch
+        return [
+            torch.cuda.get_device_name(i)
+            for i in range(torch.cuda.device_count())
+        ]
+    else:
+        return [""]

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 # Unique source of truth for the version number.
-__version__ = "0.20.0.dev202504020401"
+__version__ = "0.20.0.dev202504030357"
 @keras_hub_export("keras_hub.version")

@@ -23,7 +23,7 @@ def get_version(rel_path):
 HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
-VERSION = "0.20.0.dev202504020401"  # get_version("keras_hub/src/version_utils.py")
+VERSION = "0.20.0.dev202504030357"  # get_version("keras_hub/src/version_utils.py")
 setup(
     name="keras-hub-nightly",  # "keras-hub",

keras-hub-nightly 0.20.0.dev202504020401__tar.gz → 0.20.0.dev202504030357__tar.gz

keras-hub-nightly 0.20.0.dev202504020401tar.gz → 0.20.0.dev202504030357tar.gz