keras-hub 0.21.1.dev0__py3-none-any.whl → 0.22.0__py3-none-any.whl

keras_hub/src/models/qwen3/qwen3_decoder.py

@@ -0,0 +1,309 @@
+import keras
+from keras import ops
+
+from keras_hub.src.layers.modeling.transformer_layer_utils import (
+    compute_causal_mask,
+)
+from keras_hub.src.layers.modeling.transformer_layer_utils import (
+    merge_padding_and_attention_mask,
+)
+from keras_hub.src.models.qwen3.qwen3_attention import Qwen3Attention
+from keras_hub.src.models.qwen3.qwen3_layernorm import Qwen3LayerNorm
+from keras_hub.src.utils.keras_utils import clone_initializer
+
+
+class Qwen3TransformerDecoder(keras.layers.Layer):
+    """A Transformer decoder layer for the Qwen3 backbone.
+
+    This layer implements a Transformer decoder block that includes
+    self-attention with optional sliding window attention and a feed-forward
+    network.
+
+    Args:
+    intermediate_dim: Output dimension of the first dense layer in the
+        feed-forward network.
+    num_query_heads: Number of query attention heads.
+    num_key_value_heads: Number of key/value attention heads (for GQA).
+    rope_max_wavelength: Maximum wavelength for RoPE (Rotary Position
+        Embedding).
+    rope_scaling_factor: Scaling factor for RoPE, used for extending
+        context length.
+    activation: Activation function to use in the feed-forward network.
+    layer_norm_epsilon: Small float added to variance to avoid dividing
+        by zero in layer norm.
+    kernel_initializer: Initializer for the kernel weights.
+    dropout: Dropout rate for attention and hidden layers.
+    sliding_window_size: Size of the sliding window for attention when
+        enabled.
+    **kwargs: Additional keyword arguments to pass to the Layer.
+    """
+
+    def __init__(
+        self,
+        intermediate_dim,
+        num_query_heads,
+        num_key_value_heads,
+        head_dim,
+        rope_max_wavelength=10000,
+        rope_scaling_factor=1.0,
+        activation="silu",
+        layer_norm_epsilon=1e-5,
+        kernel_initializer="glorot_uniform",
+        dropout=0.0,
+        sliding_window_size=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.intermediate_dim = intermediate_dim
+        self.num_query_heads = num_query_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+
+        self.rope_max_wavelength = rope_max_wavelength
+        self.rope_scaling_factor = rope_scaling_factor
+
+        self.dropout = dropout
+
+        self.sliding_window_size = sliding_window_size
+
+        self.activation = keras.activations.get(activation)
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+
+        self.supports_masking = True
+
+    def build(self, decoder_sequence_shape):
+        self._decoder_sequence_shape = decoder_sequence_shape
+        self.hidden_dim = decoder_sequence_shape[-1]
+
+        # Self attention layer.
+        self._self_attention_layer = Qwen3Attention(
+            num_query_heads=self.num_query_heads,
+            num_key_value_heads=self.num_key_value_heads,
+            rope_max_wavelength=self.rope_max_wavelength,
+            head_dim=self.head_dim,
+            rope_scaling_factor=self.rope_scaling_factor,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            dropout=self.dropout,
+            sliding_window_size=self.sliding_window_size,
+            dtype=self.dtype_policy,
+            name="self_attention",
+        )
+        self._self_attention_layer.build(decoder_sequence_shape)
+
+        self._self_attention_layernorm = Qwen3LayerNorm(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="self_attention_layernorm",
+        )
+
+        self._self_attention_layernorm.build(decoder_sequence_shape)
+        self._self_attention_dropout = keras.layers.Dropout(
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="self_attention_dropout",
+        )
+
+        # Feedforward layers.
+        self._feedforward_intermediate_dense = keras.layers.Dense(
+            self.intermediate_dim,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            use_bias=False,
+            dtype=self.dtype_policy,
+            name="feedforward_intermediate_dense",
+        )
+        self._feedforward_intermediate_dense.build(decoder_sequence_shape)
+
+        self._feedforward_gate_dense = keras.layers.Dense(
+            self.intermediate_dim,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            use_bias=False,
+            dtype=self.dtype_policy,
+            name="feedforward_gate_dense",
+        )
+        self._feedforward_gate_dense.build(decoder_sequence_shape)
+
+        self._feedforward_output_dense = keras.layers.Dense(
+            self.hidden_dim,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            use_bias=False,
+            dtype=self.dtype_policy,
+            name="feedforward_output_dense",
+        )
+
+        self._feedforward_output_dense.build(
+            self._feedforward_gate_dense.compute_output_shape(
+                decoder_sequence_shape
+            )
+        )
+
+        self._feedforward_layernorm = Qwen3LayerNorm(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="feedforward_layernorm",
+        )
+        self._feedforward_layernorm.build(decoder_sequence_shape)
+
+        self.built = True
+
+    def call(
+        self,
+        decoder_sequence,
+        decoder_padding_mask=None,
+        decoder_attention_mask=None,
+        self_attention_cache=None,
+        self_attention_cache_update_index=None,
+        training=None,
+    ):
+        """Forward pass for the decoder layer.
+
+        Args:
+            decoder_sequence: Input tensor of shape [batch_size, seq_length,
+                hidden_size].
+            decoder_padding_mask: Mask tensor for padding tokens.
+            decoder_attention_mask: Additional attention mask.
+            self_attention_cache: Optional cached key and value tensors for
+                self-attention.
+            self_attention_cache_update_index: Index at which to update the
+                cache.
+            training: Boolean indicating whether in training mode.
+
+        Returns:
+            decoder_output: Output tensor after applying transformer decoder
+                block.
+            self_attention_cache: Updated cache tensors (if cache is provided).
+        """
+        self_attention_mask = self._compute_self_attention_mask(
+            decoder_sequence=decoder_sequence,
+            decoder_padding_mask=decoder_padding_mask,
+            decoder_attention_mask=decoder_attention_mask,
+            self_attention_cache=self_attention_cache,
+            self_attention_cache_update_index=self_attention_cache_update_index,
+        )
+        residual = decoder_sequence
+
+        x = self._self_attention_layernorm(decoder_sequence)
+
+        # Self attention block.
+        x = self._self_attention_layer(
+            hidden_states=x,
+            attention_mask=self_attention_mask,
+            cache=self_attention_cache,
+            cache_update_index=self_attention_cache_update_index,
+        )
+
+        if self_attention_cache is not None:
+            x, self_attention_cache = x
+
+        x = self._self_attention_dropout(x, training=training)
+
+        x = x + residual
+        residual = x
+
+        x = self._feedforward_layernorm(x)
+        gate_output = self._feedforward_gate_dense(x)
+
+        # Note that we run the activation function in full 32-bit
+        # precision since this is what `torch.nn.functional.silu`
+        # does. Internally, `torch.nn.functional.silu` converts the
+        # inputs to float32, computes SiLU, and converts the outputs
+        # back to compute dtype.
+        # CPU Kernel: https://github.com/pytorch/pytorch/blob/35c493f2cf9b623bfdc7e6b34dc1cb39690a7919/aten/src/ATen/native/cpu/Activation.cpp#L1221-L1235  # noqa: E501
+        # CUDA Kernel: https://github.com/pytorch/pytorch/blob/35c493f2cf9b623bfdc7e6b34dc1cb39690a7919/aten/src/ATen/native/cuda/ActivationSiluKernel.cu  # noqa: E501
+        gate_output = ops.cast(gate_output, "float32")
+        gate_output = self.activation(gate_output)
+        gate_output = ops.cast(gate_output, self.compute_dtype)
+
+        x = self._feedforward_intermediate_dense(x)
+
+        x = self._feedforward_output_dense(ops.multiply(x, gate_output))
+
+        decoder_output = x + residual
+
+        if self_attention_cache is not None:
+            return decoder_output, self_attention_cache
+        return decoder_output
+
+    def _compute_self_attention_mask(
+        self,
+        decoder_sequence,
+        decoder_padding_mask,
+        decoder_attention_mask,
+        self_attention_cache,
+        self_attention_cache_update_index,
+    ):
+        """Computes the self-attention mask combining causal, padding and
+        attention masks.
+
+        Args:
+            decoder_sequence: Input tensor.
+            decoder_padding_mask: Mask tensor for padding tokens.
+            decoder_attention_mask: Additional attention mask.
+            self_attention_cache: Optional cached key and value tensors.
+            self_attention_cache_update_index: Index at which to update the
+                cache.
+
+        Returns:
+            Combined attention mask tensor.
+        """
+        decoder_mask = merge_padding_and_attention_mask(
+            decoder_sequence, decoder_padding_mask, decoder_attention_mask
+        )
+        batch_size = ops.shape(decoder_sequence)[0]
+        input_length = output_length = ops.shape(decoder_sequence)[1]
+        # We need to handle a rectangular causal mask when doing cached
+        # decoding. For generative inference, `decoder_sequence` will
+        # generally be length 1, and `cache` will be the full generation length.
+        if self_attention_cache is not None:
+            input_length = ops.shape(self_attention_cache)[2]
+
+        cache_update_index = (
+            0
+            if self_attention_cache_update_index is None
+            else self_attention_cache_update_index
+        )
+
+        causal_mask = compute_causal_mask(
+            batch_size, input_length, output_length, cache_update_index
+        )
+
+        return (
+            ops.minimum(decoder_mask, causal_mask)
+            if decoder_mask is not None
+            else causal_mask
+        )
+
+    def compute_output_shape(self, decoder_sequence_shape):
+        """Computes the output shape of the layer.
+
+        Args:
+            decoder_sequence_shape: Shape of the decoder sequence input.
+
+        Returns:
+            Output shape, which is the same as the input shape.
+        """
+        return decoder_sequence_shape
+
+    def get_config(self):
+        """Returns the config of the layer.
+
+        Returns:
+            Dictionary containing the parameters used to initialize this layer.
+        """
+        config = super().get_config()
+        config.update(
+            {
+                "intermediate_dim": self.intermediate_dim,
+                "num_query_heads": self.num_query_heads,
+                "rope_max_wavelength": self.rope_max_wavelength,
+                "rope_scaling_factor": self.rope_scaling_factor,
+                "num_key_value_heads": self.num_key_value_heads,
+                "activation": keras.activations.serialize(self.activation),
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "kernel_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "dropout": self.dropout,
+            }
+        )
+        return config

keras_hub/src/models/qwen3/qwen3_layernorm.py

@@ -0,0 +1,38 @@
+import keras
+from keras import ops
+
+
+class Qwen3LayerNorm(keras.layers.Layer):
+    """A normalization layer for Qwen that implements RMS normalization."""
+
+    def __init__(self, head_dim=None, epsilon=1e-6, **kwargs):
+        super().__init__(**kwargs)
+        self.head_dim = head_dim
+        self.epsilon = epsilon
+
+    def build(self, input_shape):
+        if self.head_dim:
+            dim = self.head_dim
+        else:
+            dim = input_shape[-1]
+
+        self.scale = self.add_weight(
+            name="scale",
+            trainable=True,
+            shape=(dim,),
+            initializer="ones",
+            dtype=self.variable_dtype,
+        )
+        self.built = True
+
+    def call(self, x):
+        input_dtype = x.dtype
+        x = ops.cast(x, "float32")
+        var = ops.mean(ops.power(x, 2), axis=-1, keepdims=True)
+        x = x * ops.rsqrt(var + self.epsilon)
+        return ops.cast(x * self.scale, input_dtype)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"epsilon": self.epsilon})
+        return config

keras_hub/src/models/qwen3/qwen3_presets.py

@@ -0,0 +1,73 @@
+"""Qwen3 model preset configurations."""
+
+backbone_presets = {
+    "qwen3_0.6b_en": {
+        "metadata": {
+            "description": (
+                "28-layer Qwen3 model with 596M parameters, optimized for "
+                "efficiency and fast inference on resource-constrained devices."
+            ),
+            "params": 596049920,
+            "path": "qwen3",
+        },
+        "kaggle_handle": "kaggle://keras/qwen-3/keras/qwen3_0.6b_en/1",
+    },
+    "qwen3_1.7b_en": {
+        "metadata": {
+            "description": (
+                "28-layer Qwen3 model with 1.72B parameters, offering "
+                "a good balance between performance and resource usage."
+            ),
+            "params": 1720574976,
+            "path": "qwen3",
+        },
+        "kaggle_handle": "kaggle://keras/qwen-3/keras/qwen3_1.7b_en/1",
+    },
+    "qwen3_4b_en": {
+        "metadata": {
+            "description": (
+                "36-layer Qwen3 model with 4.02B parameters, offering improved "
+                "reasoning capabilities and better performance than smaller "
+                "variants."
+            ),
+            "params": 4022468096,
+            "path": "qwen3",
+        },
+        "kaggle_handle": "kaggle://keras/qwen-3/keras/qwen3_4b_en/1",
+    },
+    "qwen3_8b_en": {
+        "metadata": {
+            "description": (
+                "36-layer Qwen3 model with 8.19B parameters, featuring "
+                "enhanced reasoning, coding, and instruction-following "
+                "capabilities."
+            ),
+            "params": 8190735360,
+            "path": "qwen3",
+        },
+        "kaggle_handle": "kaggle://keras/qwen-3/keras/qwen3_8b_en/1",
+    },
+    "qwen3_14b_en": {
+        "metadata": {
+            "description": (
+                "40-layer Qwen3 model with 14.77B parameters, featuring "
+                "advanced reasoning, coding, and multilingual capabilities."
+            ),
+            "params": 14768307200,
+            "path": "qwen3",
+        },
+        "kaggle_handle": "kaggle://keras/qwen-3/keras/qwen3_14b_en/1",
+    },
+    "qwen3_32b_en": {
+        "metadata": {
+            "description": (
+                "64-layer Qwen3 model with 32.76B parameters, featuring "
+                "state-of-the-art performance across reasoning, coding, and "
+                "general language tasks."
+            ),
+            "params": 32762123264,
+            "path": "qwen3",
+        },
+        "kaggle_handle": "kaggle://keras/qwen-3/keras/qwen3_32b_en/1",
+    },
+}

keras_hub/src/models/qwen3/qwen3_tokenizer.py

@@ -0,0 +1,48 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.qwen3.qwen3_backbone import Qwen3Backbone
+from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+
+
+@keras_hub_export(
+    "keras_hub.models.Qwen3Tokenizer",
+)
+class Qwen3Tokenizer(BytePairTokenizer):
+    """Tokenizer for Qwen3 models.
+
+    This tokenizer implements byte-pair encoding (BPE) for Qwen3 models,
+    handling special tokens like BOS (beginning of sequence) and EOS (end of
+    sequence).
+
+    Args:
+        vocabulary: Dictionary mapping tokens to token IDs, or path to
+            vocabulary file.
+        merges: List of BPE merges, or path to merges file.
+        bos_token: Beginning of sequence token. Defaults to None.
+        eos_token: End of sequence token. Defaults to "<|endoftext|>".
+        misc_special_tokens: Set of additional special tokens. Defaults to
+            empty set.
+    """
+
+    backbone_cls = Qwen3Backbone
+
+    def __init__(
+        self,
+        vocabulary=None,
+        merges=None,
+        **kwargs,
+    ):
+        # Add EOS token
+        eos_token = "<|im_end|>"
+        self._add_special_token(eos_token, "end_token")
+
+        pad_token = "<|endoftext|>"
+        self._add_special_token(pad_token, "pad_token")
+
+        self.start_token_id = None
+        self.start_token = None
+
+        super().__init__(
+            vocabulary=vocabulary,
+            merges=merges,
+            **kwargs,
+        )

keras_hub/src/models/qwen_moe/qwen_moe_attention.py

@@ -67,6 +67,7 @@ class QwenMoeAttention(keras.layers.Layer):
         self.rope_scaling_factor = rope_scaling_factor
         self.use_sliding_window_attention = use_sliding_window_attention
         self.sliding_window_size = sliding_window_size
+        self.logit_soft_cap = None
 
     def build(self, inputs_shape):
         # Einsum variables:

keras_hub/src/models/qwen_moe/qwen_moe_presets.py

@@ -8,8 +8,8 @@ backbone_presets = {
                 "and 8 experts per MoE layer."
             ),
             "params": 14315784192,
-            "path": "qwen-1.5-moe",
+            "path": "qwen_moe",
         },
-        "kaggle_handle": "kaggle://keras/qwen-1.5-moe/Keras/qwen1.5_moe_2.7b_en/3",
+        "kaggle_handle": "kaggle://keras/qwen-1.5-moe/Keras/qwen1.5_moe_2.7b_en/4",
     },
 }

keras_hub/src/models/roformer_v2/roformer_v2_attention.py

@@ -179,8 +179,6 @@ class RoformerAttention(keras.layers.Layer):
         vw = ops.reshape(vw, (b, s, self.heads, self.head_size))
 
         qw, kw = self.rotary_embedding_layer([qw, kw])
-        if keras.__version__ < "3.6":
-            raise ("Please make sure your Keras version is >=3.6.")
         flash_attention = keras.config.is_flash_attention_enabled()
         attention_mask = ops.reshape(attention_mask, [b, 1, s, 1])
         if keras.config.backend() == "torch":

keras_hub/src/models/stable_diffusion_3/flow_match_euler_discrete_scheduler.py

@@ -38,7 +38,6 @@ class FlowMatchEulerDiscreteScheduler(layers.Layer):
         timesteps = ops.flip(timesteps, axis=0)
         sigmas = self._timestep_to_sigma(timesteps)
 
-        self.timesteps = ops.multiply(sigmas, num_train_timesteps)
         self.sigma_min = sigmas[-1]
         self.sigma_max = sigmas[0]
 
@@ -54,14 +53,24 @@ class FlowMatchEulerDiscreteScheduler(layers.Layer):
             )
         return sigma
 
+    def set_sigmas(self, num_steps):
+        timesteps = ops.linspace(
+            self._sigma_to_timestep(self.sigma_max),
+            self._sigma_to_timestep(self.sigma_min),
+            num_steps,
+        )
+        sigmas = self._timestep_to_sigma(timesteps)
+        sigmas = ops.concatenate([sigmas, ops.zeros((1,), dtype=sigmas.dtype)])
+        self.sigmas = sigmas
+
     def call(self, inputs, num_steps):
-        start = self._sigma_to_timestep(self.sigma_max)
-        end = self._sigma_to_timestep(self.sigma_min)
-        step_size = ops.divide(
-            ops.subtract(end, start), ops.subtract(num_steps, 1)
+        if not hasattr(self, "sigmas"):
+            self.set_sigmas(num_steps)
+
+        step = ops.expand_dims(
+            ops.convert_to_tensor(inputs, dtype="int32"), axis=0
         )
-        timestep = ops.add(start, ops.multiply(inputs, step_size))
-        sigma = ops.maximum(self._timestep_to_sigma(timestep), 0.0)
+        sigma = ops.take(self.sigmas, step)
         timestep = self._sigma_to_timestep(sigma)
         return sigma, timestep
 

keras_hub/src/models/stable_diffusion_3/mmdit.py

@@ -10,6 +10,63 @@ 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 standardize_data_format
 
+# TODO: Deprecate this in favor of
+# `keras.layers.RMSNormalization` once we require Keras 3.9 or later.
+if hasattr(layers, "RMSNormalization"):
+    RMSNormalization = layers.RMSNormalization
+else:
+
+    class RMSNormalization(layers.Layer):
+        """A normalization layer for MMDiT that implements RMS normalization."""
+
+        def __init__(self, axis=-1, epsilon=1e-6, **kwargs):
+            super().__init__(**kwargs)
+            self.axis = axis
+            self.epsilon = epsilon
+
+        def build(self, input_shape):
+            if isinstance(self.axis, list):
+                shape = tuple([input_shape[dim] for dim in self.axis])
+            else:
+                shape = (input_shape[self.axis],)
+                self.axis = [self.axis]
+
+            self.scale = self.add_weight(
+                name="scale", shape=shape, initializer="ones"
+            )
+
+            self.built = True
+
+        def call(self, x):
+            x = ops.cast(
+                x, keras.backend.result_type(self.compute_dtype, "float32")
+            )
+            rrms = ops.rsqrt(
+                ops.mean(ops.square(x), axis=self.axis, keepdims=True)
+                + self.epsilon
+            )
+            return (x * rrms) * ops.cast(self.scale, x.dtype)
+
+        def compute_output_shape(self, input_shape):
+            if isinstance(self.axis, int):
+                axes = [self.axis]
+            else:
+                axes = self.axis
+
+            for axis in axes:
+                if axis >= len(input_shape) or axis < -len(input_shape):
+                    raise ValueError(
+                        f"Axis {axis} is out of bounds for "
+                        f"input shape {input_shape}. "
+                        f"Received: axis={self.axis}"
+                    )
+            return input_shape
+
+        def get_config(self):
+            config = super().get_config()
+            config.update({"axis": self.axis, "epsilon": self.epsilon})
+            return config
+
 
 class AdaptiveLayerNormalization(layers.Layer):
     """Adaptive layer normalization.
@@ -402,11 +459,11 @@ def get_qk_norm(qk_norm=None, q_norm_name="q_norm", k_norm_name="k_norm"):
     if qk_norm is None:
         pass
     elif qk_norm == "rms_norm":
-        q_norm = layers.LayerNormalization(
-            epsilon=1e-6, rms_scaling=True, dtype="float32", name=q_norm_name
+        q_norm = RMSNormalization(
+            axis=-1, epsilon=1e-6, dtype="float32", name=q_norm_name
         )
-        k_norm = layers.LayerNormalization(
-            epsilon=1e-6, rms_scaling=True, dtype="float32", name=k_norm_name
+        k_norm = RMSNormalization(
+            axis=-1, epsilon=1e-6, dtype="float32", name=k_norm_name
         )
     else:
         raise NotImplementedError(