keras-hub-nightly 0.24.0.dev202511220420__py3-none-any.whl → 0.26.0.dev202601010440__py3-none-any.whl

keras_hub/src/models/dinov3/dinov3_presets.py

@@ -1,4 +1,93 @@
 """DINOV3 model preset configurations."""
 
 # Metadata for loading pretrained model weights.
-backbone_presets = {}
+backbone_presets = {
+    "dinov3_vit_small_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (small-sized model) trained on LVD-1689M "
+                "using DINOv3."
+            ),
+            "params": 21_600_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_small_lvd1689m/1",
+    },
+    "dinov3_vit_small_plus_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (small-plus-sized model) trained on "
+                "LVD-1689M using DINOv3."
+            ),
+            "params": 29_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_small_plus_lvd1689m/1",
+    },
+    "dinov3_vit_base_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (base-sized model) trained on LVD-1689M "
+                "using DINOv3."
+            ),
+            "params": 86_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_base_lvd1689m/1",
+    },
+    "dinov3_vit_large_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (large-sized model) trained on LVD-1689M "
+                "using DINOv3."
+            ),
+            "params": 300_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_large_lvd1689m/1",
+    },
+    "dinov3_vit_huge_plus_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (huge-plus-sized model) trained on "
+                "LVD-1689M using DINOv3."
+            ),
+            "params": 840_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_huge_plus_lvd1689m/1",
+    },
+    "dinov3_vit_7b_lvd1689m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (7B-sized model) trained on LVD-1689M "
+                "using DINOv3."
+            ),
+            "params": 6_700_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_7b_lvd1689m/1",
+    },
+    "dinov3_vit_large_sat493m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (large-sized model) trained on SAT-493M "
+                "using DINOv3."
+            ),
+            "params": 300_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_large_sat493m/1",
+    },
+    "dinov3_vit_7b_sat493m": {
+        "metadata": {
+            "description": (
+                "Vision Transformer (7B-sized model) trained on SAT-493M "
+                "using DINOv3."
+            ),
+            "params": 6_700_000_000,
+            "path": "dinov3",
+        },
+        "kaggle_handle": "kaggle://keras/dinov3/keras/dinov3_vit_7b_sat493m/1",
+    },
+}

keras_hub/src/models/electra/electra_backbone.py

@@ -1,10 +1,8 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.utils.keras_utils import gelu_approximate

keras_hub/src/models/esm/esm_attention.py

@@ -14,7 +14,8 @@ class ESMRotaryEmbedding(RotaryEmbedding):
         inv_freq = self.scaling_factor / (
             self.max_wavelength ** (ops.arange(0, dim, 2, dtype=x.dtype) / dim)
         )
-        t = ops.arange(x.shape[position], dtype=x.dtype)
+        # Use ops.shape for dynamic shape compatibility with TFLite
+        t = ops.arange(ops.shape(x)[position], dtype=x.dtype)
         freqs = ops.outer(t, inv_freq)
         emb = ops.concatenate((freqs, freqs), axis=-1)
 
@@ -32,11 +33,17 @@ class ESMRotaryEmbedding(RotaryEmbedding):
 
     def rotate_half(self, x):
         x1, x2 = ops.split(x, 2, -1)
-        return ops.concatenate((-x2, x1), axis=-1)
+        # Avoid `ops.concatenate` to prevent XLA compilation issues on JAX
+        # backend. Use stack + reshape approach from base RotaryEmbedding.
+        half_rot_x = ops.stack((-x2, x1), axis=-2)
+        half_rot_x = ops.reshape(half_rot_x, ops.shape(x))
+        return half_rot_x
 
     def apply_rotary_pos_emb(self, x, cos, sin):
-        cos = cos[:, : x.shape[1], :, :]
-        sin = sin[:, : x.shape[1], :, :]
+        # Use ops.shape for dynamic shape compatibility with TFLite
+        seq_len = ops.shape(x)[1]
+        cos = cos[:, :seq_len, :, :]
+        sin = sin[:, :seq_len, :, :]
 
         return (x * cos) + (self.rotate_half(x) * sin)
 

keras_hub/src/models/f_net/f_net_backbone.py

@@ -1,11 +1,9 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.f_net_encoder import FNetEncoder
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.utils.keras_utils import gelu_approximate
 

keras_hub/src/models/falcon/falcon_backbone.py

@@ -1,9 +1,7 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.falcon.falcon_transformer_decoder import (
     FalconTransformerDecoder,

keras_hub/src/models/gemma/gemma_backbone.py

@@ -1,10 +1,8 @@
 import keras
 from keras import ops
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.gemma.gemma_decoder_block import GemmaDecoderBlock
 from keras_hub.src.models.gemma.rms_normalization import RMSNormalization

keras_hub/src/models/gemma/gemma_causal_lm.py

@@ -431,3 +431,19 @@ class GemmaCausalLM(CausalLM):
         )
         per_token_loss = per_token_loss_fn(target_ids, logits)
         return per_token_loss
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        # Wrap embedding + scaling
+        backbone = self.backbone
+        inputs = keras.Input(shape=(None,), dtype="int32")
+        x = backbone.token_embedding(inputs)
+        x = x * ops.cast(ops.sqrt(backbone.hidden_dim), x.dtype)
+        pre_processor = keras.Model(inputs=inputs, outputs=x)
+
+        return {
+            "pre_block_layers": [pre_processor],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/gemma3/gemma3_backbone.py

@@ -1,10 +1,8 @@
 import keras
 from keras import ops
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.gemma.rms_normalization import RMSNormalization
 from keras_hub.src.models.gemma3.gemma3_decoder_block import Gemma3DecoderBlock

keras_hub/src/models/gemma3/gemma3_causal_lm_preprocessor.py

@@ -283,9 +283,14 @@ class Gemma3CausalLMPreprocessor(CausalLMPreprocessor):
         # is `None`.
         self.text_only_model = self.image_converter is None
 
-        self.image_placeholder = self.tokenizer.image_placeholder
-        self.start_of_image_token = self.tokenizer.start_of_image_token
-        self.end_of_image_token = self.tokenizer.end_of_image_token
+        if self.text_only_model:
+            self.image_placeholder = None
+            self.start_of_image_token = None
+            self.end_of_image_token = None
+        else:
+            self.image_placeholder = self.tokenizer.image_placeholder
+            self.start_of_image_token = self.tokenizer.start_of_image_token
+            self.end_of_image_token = self.tokenizer.end_of_image_token
 
     def build(self, input_shape):
         # Defer packer creation to `build()` so that we can be sure tokenizer

keras_hub/src/models/gemma3/gemma3_presets.py

@@ -220,4 +220,16 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://keras/medgemma/keras/medgemma_instruct_27b_text/1",
     },
+    "function_gemma_instruct_270m": {
+        "metadata": {
+            "description": (
+                "A 270M Million parameter text-only model based on Gemma 3. "
+                "This model is trained specifically for function calling "
+                "improvements."
+            ),
+            "params": 268098176,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/function-gemma/keras/function_gemma_instruct_270m/1",
+    },
 }

keras_hub/src/models/gemma3/gemma3_tokenizer.py

@@ -77,20 +77,32 @@ class Gemma3Tokenizer(SentencePieceTokenizer):
 
     backbone_cls = Gemma3Backbone
 
-    def __init__(self, proto, **kwargs):
+    def __init__(self, proto, has_vision_tokens=True, **kwargs):
         # Add special tokens.
 
+        self.has_vision_tokens = has_vision_tokens
         # The usual tokens.
         self._add_special_token("<bos>", "start_token")
         self._add_special_token("<eos>", "end_token")
         self._add_special_token("<pad>", "pad_token")
 
-        # Image placeholder token.
-        self._add_special_token("<img>", "image_placeholder")
-
-        #  Some tokens which are used in the preprocessor. We need to keep them
-        # here so that the preprocessor works with `tf.data`.
-        self._add_special_token("<start_of_image>", "start_of_image_token")
-        self._add_special_token("<end_of_image>", "end_of_image_token")
+        if has_vision_tokens:
+            # Image placeholder token.
+            self._add_special_token("<img>", "image_placeholder")
+            # Some tokens which are used in the preprocessor.
+            # We need to keep them
+            # here so that the preprocessor works with tf.data.
+            self._add_special_token("<start_of_image>", "start_of_image_token")
+            self._add_special_token("<end_of_image>", "end_of_image_token")
+        else:
+            # For text-only, skip assigning token IDs or set to -1
+            self.start_of_image_token_id = -1
+            self.image_placeholder_token_id = -1
+            self.end_of_image_token_id = -1
 
         super().__init__(proto=proto, **kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"has_vision_tokens": self.has_vision_tokens})
+        return config

keras_hub/src/models/gpt2/gpt2_backbone.py

@@ -1,10 +1,8 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.layers.modeling.transformer_decoder import TransformerDecoder
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.utils.keras_utils import gelu_approximate

keras_hub/src/models/gpt2/gpt2_causal_lm.py

@@ -420,3 +420,20 @@ class GPT2CausalLM(CausalLM):
         )
         per_token_loss = per_token_loss_fn(target_ids, logits)
         return per_token_loss
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        backbone = self.backbone
+        token_ids = keras.Input(shape=(None,), dtype="int32")
+        tokens = backbone.token_embedding(token_ids)
+        positions = backbone.position_embedding(tokens)
+        x = backbone.embeddings_add((tokens, positions))
+        x = backbone.embeddings_dropout(x)
+        pre_processor = keras.Model(inputs=token_ids, outputs=x)
+
+        return {
+            "pre_block_layers": [pre_processor],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/gpt_neo_x/gpt_neo_x_backbone.py

@@ -1,9 +1,7 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.gpt_neo_x.gpt_neo_x_decoder import GPTNeoXDecoder
 from keras_hub.src.utils.keras_utils import gelu_approximate

keras_hub/src/models/gpt_oss/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.gpt_oss.gpt_oss_backbone import GptOssBackbone
+from keras_hub.src.models.gpt_oss.gpt_oss_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, GptOssBackbone)

keras_hub/src/models/gpt_oss/gpt_oss_attention.py

@@ -0,0 +1,330 @@
+import math
+
+import keras
+from keras import ops
+
+from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
+from keras_hub.src.utils.keras_utils import clone_initializer
+
+
+class GptOssAttention(keras.layers.Layer):
+    """A cached attention layer with sliding window and sink tokens.
+
+    This layer implements the attention mechanism described in the GPT-OSS
+    paper. It includes grouped-query attention, rotary position embeddings,
+    sliding window attention, and sink tokens for improved performance on
+    long sequences.
+
+    Args:
+        num_query_heads: int. The number of query attention heads.
+        num_key_value_heads: int. The number of key and value attention
+            heads.
+        rope_max_wavelength: int. The maximum wavelength for the
+            rotary position embedding. Defaults to 10000.
+        rope_scaling_factor: float. The scaling factor for the
+            rotary position embedding. Defaults to 1.0.
+        kernel_initializer: str. The initializer for the kernel
+            weights. Defaults to "glorot_uniform".
+        sliding_window: int. The size of the sliding window.
+            Defaults to 4096.
+        dropout: float. The dropout rate. Defaults to 0.
+        head_dim: int. Head dimension for attention. If None,
+            calculated as hidden_dim // num_query_heads. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        num_query_heads,
+        num_key_value_heads,
+        rope_max_wavelength=10000,
+        rope_scaling_factor=1.0,
+        kernel_initializer="glorot_uniform",
+        sliding_window=4096,
+        dropout=0,
+        head_dim=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_query_heads = num_query_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.sliding_window = sliding_window
+        self.dropout = dropout
+        self.head_dim = head_dim
+        self.rope_max_wavelength = rope_max_wavelength
+        self.rope_scaling_factor = rope_scaling_factor
+        self.num_key_value_groups = num_query_heads // num_key_value_heads
+        self._kernel_initializer = keras.initializers.get(
+            clone_initializer(kernel_initializer)
+        )
+
+    def build(self, inputs_shape):
+        # Einsum variables:
+        # b = batch size
+        # q = query length
+        # k = key/value length
+        # m = the model's hidden_dim
+        # u = num query heads
+        # v = num key/value heads
+        # h = head dim
+        self._hidden_dim = inputs_shape[-1]
+
+        if self.head_dim is not None:
+            self._head_dim = self.head_dim
+        else:
+            self._head_dim = self._hidden_dim // self.num_query_heads
+        self._inv_norm_factor = 1.0 / math.sqrt(self._head_dim)
+
+        self._rotary_dim = (self._head_dim // 2) * 2
+
+        self.query_dense = keras.layers.EinsumDense(
+            equation="bqm,muh->bquh",
+            output_shape=(None, self.num_query_heads, self._head_dim),
+            bias_axes="uh",
+            kernel_initializer=self._kernel_initializer,
+            bias_initializer="zeros",
+            dtype=self.dtype_policy,
+            name="query",
+        )
+        self.query_dense.build(inputs_shape)
+
+        self.key_dense = keras.layers.EinsumDense(
+            equation="bkm,mvh->bkvh",
+            output_shape=(
+                None,
+                self.num_key_value_heads,
+                self._head_dim,
+            ),
+            bias_axes="vh",
+            kernel_initializer=self._kernel_initializer,
+            bias_initializer="zeros",
+            dtype=self.dtype_policy,
+            name="key",
+        )
+        self.key_dense.build(inputs_shape)
+
+        self.value_dense = keras.layers.EinsumDense(
+            equation="bkm,mvh->bkvh",
+            output_shape=(
+                None,
+                self.num_key_value_heads,
+                self._head_dim,
+            ),
+            bias_axes="vh",
+            kernel_initializer=self._kernel_initializer,
+            bias_initializer="zeros",
+            dtype=self.dtype_policy,
+            name="value",
+        )
+        self.value_dense.build(inputs_shape)
+
+        self.dropout_layer = keras.layers.Dropout(
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+        )
+
+        self.output_dense = keras.layers.EinsumDense(
+            equation="bquh,uhm->bqm",
+            output_shape=(None, self._hidden_dim),
+            bias_axes="m",
+            kernel_initializer=self._kernel_initializer,
+            bias_initializer="zeros",
+            dtype=self.dtype_policy,
+            name="attention_output",
+        )
+        self.output_dense.build(
+            (None, None, self.num_query_heads, self._head_dim)
+        )
+
+        self.rotary_embedding_layer = RotaryEmbedding(
+            max_wavelength=self.rope_max_wavelength,
+            scaling_factor=self.rope_scaling_factor,  # YaRN scaling factor
+            rope_type="yarn",
+            beta_fast=32.0,
+            beta_slow=1.0,
+            original_max_position_embeddings=4096,
+            dtype=self.dtype_policy,
+        )
+
+        self.sinks = self.add_weight(
+            shape=(self.num_query_heads,),
+            initializer="random_normal",
+            dtype=self.dtype,
+            name="sinks",
+        )
+
+        self._dot_product_equation = "bquh,bkuh->buqk"
+        self._combine_equation = "buqk,bkuh->bquh"
+
+        self.built = True
+
+    def call(
+        self,
+        hidden_states,
+        attention_mask=None,
+        cache=None,
+        cache_update_index=None,
+        training=None,
+    ):
+        start_index = (
+            cache_update_index if cache_update_index is not None else 0
+        )
+
+        query = self.query_dense(hidden_states)
+
+        # Compute RoPE for queries (only
+        #  to first _rotary_dim dimensions)
+        if self._rotary_dim < self._head_dim:
+            query_rot = query[..., : self._rotary_dim]
+            query_rot = self.rotary_embedding_layer(
+                query_rot, start_index=start_index
+            )
+            query = ops.concatenate(
+                [query_rot, query[..., self._rotary_dim :]], axis=-1
+            )
+        else:
+            query = self.rotary_embedding_layer(query, start_index=start_index)
+
+        def _compute_key_value(x):
+            key, value = self.key_dense(x), self.value_dense(x)
+            # Compute RoPE for keys (only apply to first _rotary_dim dimensions)
+            if self._rotary_dim < self._head_dim:
+                key_rot = key[..., : self._rotary_dim]
+                key_rot = self.rotary_embedding_layer(
+                    key_rot, start_index=start_index
+                )
+                key = ops.concatenate(
+                    [key_rot, key[..., self._rotary_dim :]], axis=-1
+                )
+            else:
+                key = self.rotary_embedding_layer(key, start_index=start_index)
+            return key, value
+
+        if cache is not None:
+            key_cache = cache[:, 0, ...]
+            value_cache = cache[:, 1, ...]
+            if cache_update_index is None:
+                key = key_cache
+                value = value_cache
+            else:
+                key_update, value_update = _compute_key_value(hidden_states)
+                start = [0, cache_update_index, 0, 0]
+                key = ops.slice_update(key_cache, start, key_update)
+                value = ops.slice_update(value_cache, start, value_update)
+                cache = ops.stack((key, value), axis=1)
+        else:
+            if cache_update_index is not None:
+                raise ValueError(
+                    "`cache_update_index` should not be set if `cache` is "
+                    f"`None`. Received: cache={cache}, "
+                    f"cache_update_index={cache_update_index}"
+                )
+            key, value = _compute_key_value(hidden_states)
+
+        # [batch_shape, seq_len, num_key_value_heads, head_dim]
+        # -> [batch_shape, seq_len, num_heads, head_dim]
+        key = ops.repeat(key, repeats=self.num_key_value_groups, axis=2)
+        value = ops.repeat(value, repeats=self.num_key_value_groups, axis=2)
+
+        attention_output = self._compute_attention(
+            query, key, value, attention_mask, start_index
+        )
+
+        attention_output = self.dropout_layer(
+            attention_output, training=training
+        )
+
+        attention_output = self.output_dense(attention_output)
+
+        if cache is not None:
+            return attention_output, cache
+        return attention_output
+
+    def _compute_attention(
+        self, query, key, value, attention_mask=None, start_index=0
+    ):
+        attention_scores = ops.einsum(self._dot_product_equation, query, key)
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self._inv_norm_factor, self.compute_dtype),
+        )
+
+        # Apply sliding window mask if specified
+        if self.sliding_window is not None and self.sliding_window > 0:
+            q_len = ops.shape(attention_scores)[-2]
+            kv_len = ops.shape(attention_scores)[-1]
+
+            # Query positions are offset by start_index during generation
+            q_positions = ops.arange(q_len) + start_index
+            kv_positions = ops.arange(kv_len)
+
+            # Mask true for positions outside sliding window
+            # For causal attention: mask if kv_pos < q_pos - sliding_window
+            mask = (
+                kv_positions[None, :]
+                >= q_positions[:, None] - self.sliding_window
+            )
+            if self.compute_dtype == "float32":
+                sliding_adder = ops.cast(-1e9, self.compute_dtype)
+            else:
+                sliding_adder = ops.cast(-1e4, self.compute_dtype)
+            attention_scores = ops.where(
+                mask[None, None, :, :], attention_scores, sliding_adder
+            )
+
+        if attention_mask is not None:
+            # The mask is a boolean tensor, True for positions to be masked.
+            # We add a large negative number to the masked positions.
+            # Use a large negative value for masking
+            if self.compute_dtype == "float32":
+                adder = ops.cast(-1e9, self.compute_dtype)
+            else:
+                adder = ops.cast(-1e4, self.compute_dtype)
+            attention_scores = ops.where(
+                attention_mask[:, None, :, :], attention_scores, adder
+            )
+
+        # Handle sink tokens by concatenating them to the logits.
+        b = ops.shape(attention_scores)[0]
+        q = ops.shape(attention_scores)[2]
+
+        sinks = ops.reshape(self.sinks, (1, self.num_query_heads, 1, 1))
+        sinks = ops.broadcast_to(sinks, (b, self.num_query_heads, q, 1))
+        # attention_scores shape: [b, num_heads, q, k]
+        # sinks shape: [b, num_heads, q, 1]
+        # We need to concatenate along the last dimension
+        combined_logits = ops.concatenate([attention_scores, sinks], axis=-1)
+
+        # Stabilize logits before softmax for numerical stability.
+        max_logits = ops.max(combined_logits, axis=-1, keepdims=True)
+        max_logits = ops.stop_gradient(max_logits)
+        combined_logits = combined_logits - max_logits
+
+        probs = ops.softmax(combined_logits, axis=-1)
+
+        # Remove the sink probabilities before computing the output.
+        attention_scores = probs[..., :-1]
+        attention_scores = ops.cast(attention_scores, self.compute_dtype)
+
+        attention_output = ops.einsum(
+            self._combine_equation, attention_scores, value
+        )
+
+        return attention_output
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_query_heads": self.num_query_heads,
+                "num_key_value_heads": self.num_key_value_heads,
+                "rope_max_wavelength": self.rope_max_wavelength,
+                "rope_scaling_factor": self.rope_scaling_factor,
+                "kernel_initializer": keras.initializers.serialize(
+                    self._kernel_initializer
+                ),
+                "sliding_window": self.sliding_window,
+                "dropout": self.dropout,
+                "head_dim": self.head_dim,
+            }
+        )
+        return config