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

keras_hub/models/__init__.py

@@ -340,6 +340,18 @@ from keras_hub.src.models.gpt_neo_x.gpt_neo_x_causal_lm_preprocessor import (
 from keras_hub.src.models.gpt_neo_x.gpt_neo_x_tokenizer import (
     GPTNeoXTokenizer as GPTNeoXTokenizer,
 )
+from keras_hub.src.models.gpt_oss.gpt_oss_backbone import (
+    GptOssBackbone as GptOssBackbone,
+)
+from keras_hub.src.models.gpt_oss.gpt_oss_causal_lm import (
+    GptOssCausalLM as GptOssCausalLM,
+)
+from keras_hub.src.models.gpt_oss.gpt_oss_causal_lm_preprocessor import (
+    GptOssCausalLMPreprocessor as GptOssCausalLMPreprocessor,
+)
+from keras_hub.src.models.gpt_oss.gpt_oss_tokenizer import (
+    GptOssTokenizer as GptOssTokenizer,
+)
 from keras_hub.src.models.hgnetv2.hgnetv2_backbone import (
     HGNetV2Backbone as HGNetV2Backbone,
 )

keras_hub/src/layers/modeling/reversible_embedding.py

@@ -1,281 +1,8 @@
-import inspect
-
 import keras
-from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
 
 
 @keras_hub_export("keras_hub.layers.ReversibleEmbedding")
-class ReversibleEmbedding(keras.layers.Embedding):
-    """An embedding layer which can project backwards to the input dim.
-
-    This layer is an extension of `keras.layers.Embedding` for language models.
-    This layer can be called "in reverse" with `reverse=True`, in which case the
-    layer will linearly project from `output_dim` back to `input_dim`.
-
-    By default, the reverse projection will use the transpose of the
-    `embeddings` weights to project to `input_dim` (weights are "tied"). If
-    `tie_weights=False`, the model will use a separate, trainable variable for
-    reverse projection.
-
-    This layer has no bias terms.
-
-    Args:
-        input_dim: Integer. Size of the vocabulary,
-            i.e. maximum integer index + 1.
-        output_dim: Integer. Dimension of the dense embedding.
-        tie_weights: Boolean, whether or not the matrix for embedding and
-            the matrix for the `reverse` projection should share the same
-            weights.
-        embeddings_initializer: Initializer for the `embeddings`
-            matrix (see `keras.initializers`).
-        embeddings_regularizer: Regularizer function applied to
-            the `embeddings` matrix (see `keras.regularizers`).
-        embeddings_constraint: Constraint function applied to
-            the `embeddings` matrix (see `keras.constraints`).
-        mask_zero: Boolean, whether or not the input value 0 is a special
-            "padding" value that should be masked out.
-        reverse_dtype: The dtype for the reverse projection computation.
-            Defaults to the `compute_dtype` of the layer.
-        logit_soft_cap: If `logit_soft_cap` is set and `reverse=True`, the
-            output logits will be scaled by
-            `tanh(logits / logit_soft_cap) * logit_soft_cap`. This narrows the
-            range of output logits and can improve training.
-        **kwargs: other keyword arguments passed to `keras.layers.Embedding`,
-            including `name`, `trainable`, `dtype` etc.
-
-    Call arguments:
-        inputs: The tensor inputs to the layer.
-        reverse: Boolean. If `True` the layer will perform a linear projection
-            from `output_dim` to `input_dim`, instead of a normal embedding
-            call. Default to `False`.
-
-    Example:
-    ```python
-    batch_size = 16
-    vocab_size = 100
-    hidden_dim = 32
-    seq_length = 50
-
-    # Generate random inputs.
-    token_ids = np.random.randint(vocab_size, size=(batch_size, seq_length))
-
-    embedding = keras_hub.layers.ReversibleEmbedding(vocab_size, hidden_dim)
-    # Embed tokens to shape `(batch_size, seq_length, hidden_dim)`.
-    hidden_states = embedding(token_ids)
-    # Project hidden states to shape `(batch_size, seq_length, vocab_size)`.
-    logits = embedding(hidden_states, reverse=True)
-    ```
-
-    References:
-    - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
-    - [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
-    """
-
-    def __init__(
-        self,
-        input_dim,
-        output_dim,
-        tie_weights=True,
-        embeddings_initializer="uniform",
-        embeddings_regularizer=None,
-        embeddings_constraint=None,
-        mask_zero=False,
-        reverse_dtype=None,
-        logit_soft_cap=None,
-        **kwargs,
-    ):
-        super().__init__(
-            input_dim,
-            output_dim,
-            embeddings_initializer=embeddings_initializer,
-            embeddings_regularizer=embeddings_regularizer,
-            embeddings_constraint=embeddings_constraint,
-            mask_zero=mask_zero,
-            **kwargs,
-        )
-        self.tie_weights = tie_weights
-        self.reverse_dtype = reverse_dtype
-        self.logit_soft_cap = logit_soft_cap
-
-    def build(self, inputs_shape=None):
-        super().build(inputs_shape)
-        if (
-            not self.tie_weights
-            and getattr(self, "quantization_mode", None) != "int8"
-        ):
-            self.reverse_embeddings = self.add_weight(
-                name="reverse_embeddings",
-                shape=(self.output_dim, self.input_dim),
-                initializer=self.embeddings_initializer,
-                dtype=self.dtype,
-            )
-
-    def call(self, inputs, reverse=False):
-        if reverse:
-            if self.tie_weights:
-                kernel = ops.transpose(ops.convert_to_tensor(self.embeddings))
-            else:
-                kernel = self.reverse_embeddings
-            if self.reverse_dtype is not None:
-                inputs = ops.cast(inputs, self.reverse_dtype)
-                kernel = ops.cast(kernel, self.reverse_dtype)
-            logits = ops.matmul(inputs, kernel)
-            # Optionally soft-cap logits.
-            if self.logit_soft_cap is not None:
-                soft_cap = self.logit_soft_cap
-                logits = ops.tanh(logits / soft_cap) * soft_cap
-            return logits
-
-        return super().call(inputs)
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "tie_weights": self.tie_weights,
-                "reverse_dtype": self.reverse_dtype,
-                "logit_soft_cap": self.logit_soft_cap,
-            }
-        )
-        return config
-
-    def save_own_variables(self, store):
-        if not self.built:
-            return
-        super().save_own_variables(store)
-        target_variables = []
-        if not self.tie_weights:
-            # Store the reverse embedding weights as the last weights.
-            target_variables.append(self.reverse_embeddings)
-            if getattr(self, "quantization_mode", None) == "int8":
-                target_variables.append(self.reverse_embeddings_scale)
-            for i, variable in enumerate(target_variables, start=len(store)):
-                store[str(i)] = variable
-
-    def load_own_variables(self, store):
-        if not self.built:
-            self.build()
-        super().load_own_variables(store)
-        if not self.tie_weights:
-            # Last weights in the stores are the reverse embedding weights.
-            target_variables = [self.reverse_embeddings]
-            if getattr(self, "quantization_mode", None) == "int8":
-                target_variables.append(self.reverse_embeddings_scale)
-            for i, variable in enumerate(
-                target_variables, start=len(store) - len(target_variables)
-            ):
-                variable.assign(store[str(i)])
-
-    def compute_output_spec(self, inputs, reverse=False):
-        output_shape = list(inputs.shape)
-        if reverse:
-            output_shape[-1] = self.input_dim
-        else:
-            output_shape += [self.output_dim]
-        return keras.KerasTensor(output_shape, dtype=self.compute_dtype)
-
-    # Quantization-related (int8) methods
-
-    def quantized_call(self, inputs, reverse=False):
-        # TODO (hongyu): This function could be removed once we add `*args` and
-        # `**kwargs` for `Embedding.quantized_call`
-        if self.quantization_mode == "int8":
-            return self._int8_call(inputs, reverse=reverse)
-        else:
-            self._quantization_mode_error(self.quantization_mode)
-
-    def _int8_build(self, embeddings_shape=None):
-        if (
-            "embeddings_shape"
-            in inspect.signature(super()._int8_build).parameters
-        ):
-            if embeddings_shape is None:
-                embeddings_shape = (self.input_dim, self.output_dim)
-            super()._int8_build(embeddings_shape=embeddings_shape)
-        else:
-            # Backward compatibility for older versions of Keras.
-            super()._int8_build()
-        self.inputs_quantizer = keras.quantizers.AbsMaxQuantizer(axis=-1)
-        if not self.tie_weights:
-            self.reverse_embeddings = self.add_weight(
-                name="reverse_embeddings",
-                shape=(self.output_dim, self.input_dim),
-                initializer="zeros",
-                dtype="int8",
-                trainable=False,
-            )
-            self.reverse_embeddings_scale = self.add_weight(
-                name="reverse_embeddings_scale",
-                shape=(self.input_dim,),
-                initializer="ones",
-                trainable=False,
-            )
-        self._is_quantized = True
-
-    def _int8_call(self, inputs, reverse=False):
-        if reverse:
-            if self.tie_weights:
-                kernel = ops.transpose(self._embeddings)
-                scale = ops.transpose(self.embeddings_scale)
-            else:
-                kernel = self.reverse_embeddings
-                scale = self.reverse_embeddings_scale
-            inputs, inputs_scale = self.inputs_quantizer(inputs)
-            logits = ops.matmul(inputs, kernel)
-            # De-scale outputs
-            logits = ops.cast(logits, self.compute_dtype)
-            logits = ops.divide(logits, ops.multiply(inputs_scale, scale))
-            # Optionally soft-cap logits.
-            if self.logit_soft_cap is not None:
-                soft_cap = self.logit_soft_cap
-                logits = ops.tanh(logits / soft_cap) * soft_cap
-            return logits
-
-        return super()._int8_call(inputs)
-
-    def quantize(self, mode, type_check=True, config=None):
-        del config
-        if type_check and type(self) is not ReversibleEmbedding:
-            raise self._not_implemented_error(self.quantize)
-
-        def abs_max_quantize(inputs, axis):
-            return keras.quantizers.abs_max_quantize(
-                inputs, axis=axis, to_numpy=True
-            )
-
-        if mode != "int8":
-            raise NotImplementedError(
-                "Invalid quantization mode. Expected 'int8'. "
-                f"Received: quantization_mode={mode}"
-            )
-
-        embeddings_shape = (self.input_dim, self.output_dim)
-        if mode == "int8":
-            embeddings, embeddings_scale = abs_max_quantize(
-                self._embeddings, axis=-1
-            )
-            embeddings_scale = ops.squeeze(embeddings_scale, axis=-1)
-            del self._embeddings
-            if not self.tie_weights:
-                reverse_embeddings, reverse_embeddings_scale = abs_max_quantize(
-                    self.reverse_embeddings, axis=0
-                )
-                reverse_embeddings_scale = ops.squeeze(
-                    reverse_embeddings_scale, axis=0
-                )
-                del self.reverse_embeddings
-        self.quantized_build(embeddings_shape, mode)
-        if mode == "int8":
-            self._embeddings.assign(embeddings)
-            self.embeddings_scale.assign(embeddings_scale)
-            if not self.tie_weights:
-                self.reverse_embeddings.assign(reverse_embeddings)
-                self.reverse_embeddings_scale.assign(reverse_embeddings_scale)
-
-        if self.dtype_policy.quantization_mode is None:
-            policy = keras.dtype_policies.get(
-                f"{mode}_from_{self.dtype_policy.name}"
-            )
-            self.dtype_policy = policy
+class ReversibleEmbedding(keras.layers.ReversibleEmbedding):
+    pass

keras_hub/src/layers/modeling/rotary_embedding.py

@@ -1,4 +1,5 @@
 import keras
+import numpy as np
 from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
@@ -25,6 +26,17 @@ class RotaryEmbedding(keras.layers.Layer):
             curves.
         scaling_factor: float. The scaling factor used to scale positions of
             the tokens.
+        rope_type: str. The type of RoPE scaling to apply. Supported types:
+            "linear", "dynamic", "yarn". Defaults to "linear".
+        beta_fast: float. Beta fast parameter for YaRN scaling. Only used
+            when rope_type="yarn". Defaults to 32.0.
+        beta_slow: float. Beta slow parameter for YaRN scaling. Only used
+            when rope_type="yarn". Defaults to 1.0.
+        original_max_position_embeddings: int. Original maximum position
+            embeddings for YaRN scaling. Only used when rope_type="yarn".
+            Defaults to 4096.
+        truncate: bool. Whether to apply truncation for YaRN scaling. Only used
+            when rope_type="yarn". Defaults to False.
         sequence_axis: int. Sequence axis in the input tensor.
         feature_axis: int. Feature axis in the input tensor.
         **kwargs: other keyword arguments passed to `keras.layers.Layer`,
@@ -69,6 +81,11 @@ class RotaryEmbedding(keras.layers.Layer):
         self,
         max_wavelength=10000,
         scaling_factor=1.0,
+        rope_type="linear",
+        beta_fast=32.0,
+        beta_slow=1.0,
+        original_max_position_embeddings=4096,
+        truncate=False,
         sequence_axis=1,
         feature_axis=-1,
         **kwargs,
@@ -78,24 +95,70 @@ class RotaryEmbedding(keras.layers.Layer):
         self.sequence_axis = sequence_axis
         self.feature_axis = feature_axis
         self.scaling_factor = scaling_factor
+        self.rope_type = rope_type
+
+        # YaRN-specific parameters (only used when rope_type="yarn")
+        self.beta_fast = beta_fast
+        self.beta_slow = beta_slow
+        self.original_max_position_embeddings = original_max_position_embeddings
+        self.truncate = truncate
         self.built = True
 
+    def _normalize_axes(self, input_shape):
+        """Normalize and validate axis indices for the given input shape."""
+        rank = len(input_shape)
+
+        # Normalize negative indices
+        sequence_axis = self.sequence_axis
+        feature_axis = self.feature_axis
+
+        if sequence_axis < 0:
+            sequence_axis += rank
+        if feature_axis < 0:
+            feature_axis += rank
+
+        if sequence_axis < 0 or sequence_axis >= rank:
+            raise ValueError(
+                f"sequence_axis {self.sequence_axis} "
+                f"is out of range for input with rank {rank}"
+            )
+        if feature_axis < 0 or feature_axis >= rank:
+            raise ValueError(
+                f"feature_axis {self.feature_axis} "
+                f"is out of range for input with rank {rank}"
+            )
+        if sequence_axis == feature_axis:
+            raise ValueError("sequence_axis and feature_axis must be different")
+
+        return sequence_axis, feature_axis
+
+    def _validate_rotary_dimension(self, rotary_dim):
+        if rotary_dim % 2 != 0:
+            raise ValueError(
+                f"Rotary dimension must be even, got {rotary_dim}."
+                "The rotary embedding splits the feature dimension "
+                "into two halves. Consider using a different feature "
+                "dimension or padding."
+            )
+
     def call(self, inputs, start_index=0, positions=None):
+        input_shape = ops.shape(inputs)
+        sequence_axis, feature_axis = self._normalize_axes(input_shape)
+
+        rotary_dim = input_shape[feature_axis]
+        self._validate_rotary_dimension(rotary_dim)
+
         # Take care of unbatched `positions`.
         if positions is not None:
             if len(ops.shape(positions)) == 1:
                 positions = ops.expand_dims(positions, axis=0)
 
-        inputs = ops.moveaxis(
-            inputs, (self.feature_axis, self.sequence_axis), (-1, 1)
-        )
+        inputs = ops.moveaxis(inputs, (feature_axis, sequence_axis), (-1, 1))
         cos_emb, sin_emb = self._compute_cos_sin_embedding(
             inputs, start_index, positions
         )
         output = self._apply_rotary_pos_emb(inputs, cos_emb, sin_emb)
-        return ops.moveaxis(
-            output, (-1, 1), (self.feature_axis, self.sequence_axis)
-        )
+        return ops.moveaxis(output, (-1, 1), (feature_axis, sequence_axis))
 
     def _apply_rotary_pos_emb(self, tensor, cos_emb, sin_emb):
         x1, x2 = ops.split(tensor, 2, axis=-1)
@@ -113,19 +176,35 @@ class RotaryEmbedding(keras.layers.Layer):
         return positions + ops.cast(start_index, dtype="float32")
 
     def _compute_cos_sin_embedding(self, inputs, start_index=0, positions=None):
+        """Compute cos & sin RoPE embeddings with optional YaRN scaling.
+        Uses tensor ops only to remain JIT/backends friendly.
+        """
         batch_axis = 0
-        feature_axis = len(inputs.shape) - 1
         sequence_axis = 1
+        feature_axis = len(inputs.shape) - 1
 
         rotary_dim = ops.shape(inputs)[feature_axis]
         inverse_freq = self._get_inverse_freq(rotary_dim)
 
         if positions is None:
             positions = self._compute_positions(inputs, start_index)
-            positions = ops.expand_dims(positions, axis=batch_axis)
+            positions = ops.expand_dims(
+                positions, axis=batch_axis
+            )  # shape (1, seq_len)
         else:
             positions = ops.cast(positions, "float32")
-        positions = positions / ops.cast(self.scaling_factor, "float32")
+            if len(ops.shape(positions)) == 1:
+                positions = ops.expand_dims(positions, axis=batch_axis)
+
+        if (
+            self.rope_type == "yarn"
+            and self.truncate
+            and self.original_max_position_embeddings is not None
+        ):
+            positions = ops.minimum(
+                positions,
+                ops.cast(self.original_max_position_embeddings, "float32"),
+            )
 
         freq = ops.einsum("bi,j->bij", positions, inverse_freq)
 
@@ -140,15 +219,103 @@ class RotaryEmbedding(keras.layers.Layer):
 
         cos_emb = ops.cast(ops.cos(embedding), self.compute_dtype)
         sin_emb = ops.cast(ops.sin(embedding), self.compute_dtype)
+
+        if self.rope_type == "yarn":
+            # YaRN temperature scaling
+            factor = ops.add(
+                ops.multiply(
+                    ops.cast(0.1, self.compute_dtype),
+                    ops.log(ops.cast(self.scaling_factor, self.compute_dtype)),
+                ),
+                ops.cast(1.0, self.compute_dtype),
+            )
+            cos_emb = cos_emb * factor
+            sin_emb = sin_emb * factor
         return cos_emb, sin_emb
 
     def _get_inverse_freq(self, rotary_dim):
-        freq_range = ops.divide(
-            ops.arange(0, rotary_dim, 2, dtype="float32"),
-            ops.cast(rotary_dim, "float32"),
+        """Return inverse frequencies."""
+        idx = ops.arange(0, rotary_dim, 2, dtype="float32")
+        denom = ops.cast(rotary_dim, "float32")
+        freq_range = idx / denom
+        inv = ops.power(ops.cast(self.max_wavelength, "float32"), -freq_range)
+
+        if self.rope_type == "linear":
+            return inv / ops.cast(self.scaling_factor, "float32")
+        elif self.rope_type == "dynamic":
+            exponent = ops.cast(rotary_dim, "float32") / ops.cast(
+                max(1, rotary_dim - 2), "float32"
+            )
+            return inv / ops.power(
+                ops.cast(self.scaling_factor, "float32"), exponent
+            )
+        elif self.rope_type == "yarn":
+            return self._get_yarn_inverse_freq(rotary_dim)
+        else:
+            return inv
+
+    def _get_yarn_inverse_freq(self, rotary_dim):
+        # Get the base (rope_theta equivalent) from max_wavelength
+        base = ops.cast(self.max_wavelength, "float32")
+
+        # Compute base frequencies: base ** (idx / dim)
+        idx = ops.arange(0, rotary_dim, 2, dtype="float32")
+        pos_freqs = ops.power(base, idx / ops.cast(rotary_dim, "float32"))
+
+        # Compute interpolation and extrapolation frequencies
+        inv_freq_extrapolation = 1.0 / pos_freqs
+        inv_freq_interpolation = 1.0 / (
+            ops.cast(self.scaling_factor, "float32") * pos_freqs
         )
-        inverse_freq = 1.0 / (self.max_wavelength**freq_range)
-        return inverse_freq
+
+        # Find correction range
+        beta_fast = ops.cast(self.beta_fast, "float32")
+        beta_slow = ops.cast(self.beta_slow, "float32")
+
+        # Find correction dimensions for beta_fast and beta_slow
+        def find_correction_dim_tensor(num_rotations, dim):
+            max_pos = ops.cast(self.original_max_position_embeddings, "float32")
+            return (dim * ops.log(max_pos / (num_rotations * 2 * np.pi))) / (
+                2 * ops.log(base)
+            )
+
+        low = find_correction_dim_tensor(
+            beta_fast, ops.cast(rotary_dim, "float32")
+        )
+        high = find_correction_dim_tensor(
+            beta_slow, ops.cast(rotary_dim, "float32")
+        )
+
+        # Apply truncation if specified
+        if self.truncate:
+            low = ops.floor(low)
+            high = ops.ceil(high)
+
+        # Clamp to valid range
+        low = ops.maximum(low, ops.cast(0, "float32"))
+        high = ops.minimum(high, ops.cast(rotary_dim // 2 - 1, "float32"))
+
+        # Linear ramp function
+        dim_half = rotary_dim // 2
+        idx_half = ops.arange(0, dim_half, dtype="float32")
+
+        # Prevent singularity
+        diff = high - low
+        diff = ops.maximum(diff, ops.cast(0.001, "float32"))
+
+        linear_func = (idx_half - low) / diff
+        ramp_func = ops.clip(linear_func, 0, 1)
+
+        # Apply the ramp to get extrapolation factor
+        inv_freq_extrapolation_factor = 1 - ramp_func
+
+        # Combine interpolation and extrapolation
+        scaled_inverse_freq = (
+            inv_freq_interpolation * (1 - inv_freq_extrapolation_factor)
+            + inv_freq_extrapolation * inv_freq_extrapolation_factor
+        )
+
+        return scaled_inverse_freq
 
     def get_config(self):
         config = super().get_config()
@@ -156,6 +323,13 @@ class RotaryEmbedding(keras.layers.Layer):
             {
                 "max_wavelength": self.max_wavelength,
                 "scaling_factor": self.scaling_factor,
+                "rope_type": self.rope_type,
+                "beta_fast": self.beta_fast,
+                "beta_slow": self.beta_slow,
+                "original_max_position_embeddings": (
+                    self.original_max_position_embeddings
+                ),
+                "truncate": self.truncate,
                 "sequence_axis": self.sequence_axis,
                 "feature_axis": self.feature_axis,
             }

keras_hub/src/layers/modeling/token_and_position_embedding.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.utils.keras_utils import clone_initializer
 
 

keras_hub/src/models/albert/albert_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/bart/bart_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.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_hub.src.models.backbone import Backbone

keras_hub/src/models/bert/bert_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/bloom/bloom_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.bloom.bloom_decoder import BloomDecoder
 

keras_hub/src/models/causal_lm.py

@@ -196,7 +196,7 @@ class CausalLM(Task):
 
                 # Create an explicit tuple of all variable state.
                 state = (
-                    self.sampler.variables,
+                    [v.value for v in self.sampler.variables],
                     # Use the explicit variable.value to preserve the
                     # sharding spec of distribution.
                     [v.value for v in self.trainable_variables],
@@ -429,3 +429,25 @@ class CausalLM(Task):
         super()._post_quantize(mode, **kwargs)
         # Reset the compiled generate function.
         self.generate_function = None
+
+    def get_quantization_layer_structure(self, mode):
+        if mode != "gptq":
+            return None
+
+        backbone = self.backbone
+        # Check for standard backbone structure.
+        if not hasattr(backbone, "transformer_layers"):
+            return None
+
+        # Check for embedding.
+        embedding = getattr(backbone, "token_embedding", None)
+        if embedding is None:
+            embedding = getattr(backbone, "embedding", None)
+
+        if embedding is None:
+            return None
+
+        return {
+            "pre_block_layers": [embedding],
+            "sequential_blocks": backbone.transformer_layers,
+        }

keras_hub/src/models/deberta_v3/deberta_v3_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.deberta_v3.disentangled_attention_encoder import (
     DisentangledAttentionEncoder,