keras-nightly 3.12.0.dev2025100503__py3-none-any.whl → 3.14.0.dev2026011604__py3-none-any.whl

keras/src/quantizers/quantization_config.py

@@ -0,0 +1,246 @@
+from keras.src.api_export import keras_export
+from keras.src.dtype_policies import QUANTIZATION_MODES
+from keras.src.saving import serialization_lib
+
+
+@keras_export("keras.quantizers.QuantizationConfig")
+class QuantizationConfig:
+    """Base class for quantization configs.
+
+    Subclasses must implement the `mode` property and the `get_config` and
+    `from_config` class methods.
+
+    Args:
+        weight_quantizer: Quantizer for weights.
+        activation_quantizer: Quantizer for activations.
+    """
+
+    def __init__(self, weight_quantizer=None, activation_quantizer=None):
+        self.weight_quantizer = weight_quantizer
+        self.activation_quantizer = activation_quantizer
+
+    @property
+    def mode(self):
+        raise NotImplementedError(
+            "Subclasses must implement this property. Do not instantiate "
+            "QuantizationConfig directly."
+        )
+
+    def get_config(self):
+        return {
+            "weight_quantizer": serialization_lib.serialize_keras_object(
+                self.weight_quantizer
+            ),
+            "activation_quantizer": serialization_lib.serialize_keras_object(
+                self.activation_quantizer
+            ),
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        weight_quantizer = serialization_lib.deserialize_keras_object(
+            config.get("weight_quantizer")
+        )
+        activation_quantizer = serialization_lib.deserialize_keras_object(
+            config.get("activation_quantizer")
+        )
+        return cls(
+            weight_quantizer=weight_quantizer,
+            activation_quantizer=activation_quantizer,
+        )
+
+    @staticmethod
+    def weight_quantizer_or_default(config, default):
+        if config is not None and config.weight_quantizer is not None:
+            return config.weight_quantizer
+        return default
+
+    @staticmethod
+    def activation_quantizer_or_default(config, default):
+        if config is not None:
+            return config.activation_quantizer
+        return default
+
+
+@keras_export("keras.quantizers.Int8QuantizationConfig")
+class Int8QuantizationConfig(QuantizationConfig):
+    """Int8 quantization config.
+
+    Args:
+        weight_quantizer: Quantizer for weights.
+        activation_quantizer: Quantizer for activations. If "default", uses
+            AbsMaxQuantizer with axis=-1.
+    """
+
+    def __init__(self, weight_quantizer=None, activation_quantizer="default"):
+        from keras.src.quantizers.quantizers import AbsMaxQuantizer
+
+        if activation_quantizer == "default":
+            activation_quantizer = AbsMaxQuantizer()
+        super().__init__(weight_quantizer, activation_quantizer)
+        if self.weight_quantizer is not None:
+            if self.weight_quantizer.output_dtype != "int8":
+                raise ValueError(
+                    "Int8QuantizationConfig requires a weight_quantizer "
+                    "with output_dtype='int8'. Received: "
+                    f"output_dtype={self.weight_quantizer.output_dtype}"
+                )
+
+    @property
+    def mode(self):
+        return "int8"
+
+
+@keras_export("keras.quantizers.Int4QuantizationConfig")
+class Int4QuantizationConfig(QuantizationConfig):
+    """Int4 quantization config.
+
+    Args:
+        weight_quantizer: Quantizer for weights.
+        activation_quantizer: Quantizer for activations. If "default", uses
+            AbsMaxQuantizer with axis=-1.
+    """
+
+    def __init__(self, weight_quantizer=None, activation_quantizer="default"):
+        from keras.src.quantizers.quantizers import AbsMaxQuantizer
+
+        if activation_quantizer == "default":
+            activation_quantizer = AbsMaxQuantizer()
+        super().__init__(weight_quantizer, activation_quantizer)
+        if self.weight_quantizer is not None:
+            if self.weight_quantizer.value_range != (-8, 7):
+                raise ValueError(
+                    "Int4QuantizationConfig requires a weight_quantizer "
+                    "with value_range=(-8, 7). Received: "
+                    f"value_range={self.weight_quantizer.value_range}"
+                )
+
+            if self.weight_quantizer.output_dtype != "int8":
+                raise ValueError(
+                    "Int4QuantizationConfig requires a weight_quantizer "
+                    "with output_dtype='int8'. Received: "
+                    f"output_dtype={self.weight_quantizer.output_dtype}"
+                )
+
+    @property
+    def mode(self):
+        return "int4"
+
+
+@keras_export("keras.quantizers.Float8QuantizationConfig")
+class Float8QuantizationConfig(QuantizationConfig):
+    """FP8 quantization config.
+
+    FP8 mixed-precision training does not support user defined quantizers.
+    This config is only used to indicate that FP8 mixed-precision training
+    should be used.
+    """
+
+    def __init__(self):
+        super().__init__(None, None)
+
+    @property
+    def mode(self):
+        return "float8"
+
+    def get_config(self):
+        return {}
+
+    @classmethod
+    def from_config(cls, config):
+        return cls()
+
+
+def validate_and_resolve_config(mode, config):
+    """Validate and resolve quantization config.
+
+    This function validates the quantization config and resolves the mode.
+    If mode is not provided, it is inferred from the config.
+    If config is not provided, a default config is inferred from the mode.
+
+    Args:
+        mode: Quantization mode.
+        config: Quantization config.
+    """
+    # 1. Backwards Compatibility: Handle string shortcuts.
+    if isinstance(config, str):
+        mode = config
+        config = None
+
+    _validate_mode(mode)
+
+    # 2. Resolve "mode" into a Config object.
+    if config is None:
+        if mode == "int8":
+            config = Int8QuantizationConfig()
+        elif mode == "int4":
+            config = Int4QuantizationConfig()
+        elif mode == "float8":
+            config = Float8QuantizationConfig()
+        elif mode == "gptq":
+            raise ValueError(
+                "For GPTQ, you must pass a `GPTQConfig` object in the "
+                "`config` argument."
+            )
+        elif mode == "awq":
+            raise ValueError(
+                "For AWQ, you must pass an `AWQConfig` object in the "
+                "`config` argument."
+            )
+        else:
+            if mode is not None:
+                raise ValueError(
+                    f"Invalid quantization mode. Received: mode={mode}"
+                )
+            raise ValueError(
+                "You must provide either `mode` or `config` to `quantize`."
+            )
+    else:
+        if not isinstance(config, QuantizationConfig):
+            raise ValueError(
+                "Argument `config` must be an instance of "
+                "`QuantizationConfig`. "
+                f"Received: config={config} (of type {type(config)})"
+            )
+
+    # 3. Validation: Prevent contradictions.
+    if mode is not None and config.mode != mode:
+        raise ValueError(
+            f"Contradictory arguments: mode='{mode}' but "
+            f"config.mode='{config.mode}'"
+        )
+
+    # Ensure mode is consistent.
+    mode = config.mode
+
+    # Ensure the mode derived from the config is valid.
+    _validate_mode(mode)
+
+    if mode == "gptq":
+        from keras.src.quantizers.gptq_config import GPTQConfig
+
+        if not isinstance(config, GPTQConfig):
+            raise ValueError(
+                "Mode 'gptq' requires a valid `config` argument of type "
+                f"`GPTQConfig`. Received: {type(config)}"
+            )
+
+    if mode == "awq":
+        from keras.src.quantizers.awq_config import AWQConfig
+
+        if not isinstance(config, AWQConfig):
+            raise ValueError(
+                "Mode 'awq' requires a valid `config` argument of type "
+                f"`AWQConfig`. Received: {type(config)}"
+            )
+
+    return config
+
+
+def _validate_mode(mode):
+    """Validates quantization mode."""
+    if mode is not None and mode not in QUANTIZATION_MODES:
+        raise ValueError(
+            "Invalid quantization mode. "
+            f"Expected one of {QUANTIZATION_MODES}. Received: mode={mode}"
+        )

keras/src/quantizers/quantizers.py

@@ -73,6 +73,23 @@ def abs_max_quantize(
     epsilon=backend.epsilon(),
     to_numpy=False,
 ):
+    """
+    Quantizes the input tensor using the absolute maximum quantization scheme.
+
+    Args:
+        inputs: Input tensor to quantize.
+        axis: Axis along which to compute the quantization range.
+        value_range: Tuple of the minimum and maximum values of the quantization
+            range.
+        dtype: Data type of the quantized output.
+        epsilon: Small value to avoid division by zero.
+        to_numpy: Whether to perform the quantization in numpy. This performs
+            the computation on the host CPU and can be useful for saving memory
+            on the device. If False, the computation is performed on the device.
+
+    Returns:
+        A tuple of the quantized tensor and the scale.
+    """
     if to_numpy:
         # Save memory on the device using numpy
         original_dtype = backend.standardize_dtype(inputs.dtype)
@@ -105,31 +122,69 @@ def abs_max_quantize(
 class AbsMaxQuantizer(Quantizer):
     def __init__(
         self,
-        axis,
+        axis=None,  # Deprecated, provide axis in __call__ instead.
         value_range=(-127, 127),
         epsilon=backend.epsilon(),
         output_dtype="int8",
     ):
         Quantizer.__init__(self, output_dtype=output_dtype)
-        if isinstance(axis, int):
-            axis = (axis,)
-        self.axis = tuple(axis)
+        if axis is not None:
+            if isinstance(axis, int):
+                axis = (axis,)
+            self.axis = tuple(axis)
+        else:
+            self.axis = None
         self.value_range = value_range
         self.epsilon = epsilon
+        if output_dtype == "int8":
+            if value_range[0] < -128 or value_range[1] > 127:
+                raise ValueError(
+                    f"Quantizer with output_dtype='int8' requires value_range "
+                    f"to be within the interval [-128, 127]. Received: "
+                    f"value_range={value_range}"
+                )
 
-    def __call__(self, x):
+    def __call__(self, x, axis=None, to_numpy=False):
+        """
+        Quantizes the input tensor.
+
+        Args:
+            x: Input tensor to quantize.
+            axis: Axis along which to compute the quantization range. If None,
+                uses the axis specified in the constructor. If None and no axis
+                was specified in the constructor, defaults to -1.
+            to_numpy: Whether to perform the quantization in numpy. This
+                performs the computation on the host CPU and can be useful for
+                saving memory on the device. If False, the computation is
+                performed on the device.
+
+        Returns:
+            A tuple of the quantized tensor and the scale.
+        """
+        if axis is None:
+            axis = self.axis
+        if axis is None:
+            # Default to -1 if no axis is specified
+            axis = -1
         quantized_x, scale = abs_max_quantize(
-            x, self.axis, self.value_range, self.output_dtype, self.epsilon
+            x,
+            axis,
+            self.value_range,
+            self.output_dtype,
+            self.epsilon,
+            to_numpy,
         )
         return quantized_x, scale
 
     def get_config(self):
-        return {
-            "axis": self.axis,
+        config = {
             "value_range": self.value_range,
             "epsilon": self.epsilon,
             "output_dtype": self.output_dtype,
         }
+        if self.axis is not None:
+            config["axis"] = self.axis
+        return config
 
 
 def adjust_and_nudge(min_range, max_range, num_bits, narrow_range):
@@ -281,7 +336,7 @@ def fake_quant_with_min_max_vars(
             ops.add(ops.multiply(-nudged_min, inv_scale), 0.5)
         )
         x_clamped = ops.clip(
-            x, ops.cast(nudged_min, x.dtype), ops.cast(nudged_max, x.dtype)
+            ops.cast(x, nudged_min.dtype), nudged_min, nudged_max
         )
         x_clamped_shifted = ops.subtract(x_clamped, nudged_min)
         result = ops.multiply(
@@ -318,6 +373,7 @@ def fake_quant_with_min_max_vars(
                 grad_min = ops.sum(grad_min, axis=axes)
             else:
                 grad_min = ops.sum(grad_min)
+            grad_min = ops.reshape(grad_min, ops.shape(min_val))
 
             # Gradient for max_val
             # When x is clipped to max, the gradient flows to max_val
@@ -327,6 +383,7 @@ def fake_quant_with_min_max_vars(
                 grad_max = ops.sum(grad_max, axis=axes)
             else:
                 grad_max = ops.sum(grad_max)
+            grad_max = ops.reshape(grad_max, ops.shape(max_val))
 
             return dx, grad_min, grad_max
 
@@ -596,11 +653,14 @@ def unpack_int4(packed, orig_len, axis=0, dtype="int8"):
         )
 
     def to_signed(x):
-        """Converts unpacked nibbles [0, 15] to signed int4 [-8, 7]."""
+        """Converts unpacked nibbles [0, 15] to signed int4 [-8, 7].
+
+        Uses a branchless XOR approach: (x ^ 8) - 8
+        This maps: 0->0, 1->1, ..., 7->7, 8->-8, 9->-7, ..., 15->-1
+        """
         dtype_x = backend.standardize_dtype(x.dtype)
         eight = ops.cast(8, dtype_x)
-        sixteen = ops.cast(16, dtype_x)
-        return ops.where(x < eight, x, x - sixteen)
+        return ops.subtract(ops.bitwise_xor(x, eight), eight)
 
     rank = getattr(packed.shape, "rank", None) or len(packed.shape)
     if axis < 0:
@@ -691,7 +751,7 @@ class GPTQQuantizer(Quantizer):
         self.zero = None
         self.maxq = None
 
-    def find_params(self, input_tensor, weight=True):
+    def find_params(self, input_tensor):
         """Finds quantization parameters (scale and zero) for a given tensor."""
         self.scale, self.zero, self.maxq = compute_quantization_parameters(
             input_tensor,
@@ -699,7 +759,6 @@ class GPTQQuantizer(Quantizer):
             symmetric=self.symmetric,
             per_channel=self.per_channel,
             group_size=self.group_size,
-            weight=weight,
             compute_dtype=self.compute_dtype,
         )
         return self.scale, self.zero, self.maxq
@@ -736,98 +795,105 @@ def compute_quantization_parameters(
     symmetric=False,
     per_channel=False,
     group_size=-1,
-    weight=False,
     compute_dtype="float32",
 ):
     """
-    Computes the scale and zero-point for quantization.
+    Computes the scale and zero-point for quantizing weight tensors.
 
     This function calculates the scale and zero-point required for quantizing
-    a given tensor `x` based on the specified parameters. It supports grouped,
-    per-channel, per-tensor, symmetric, and asymmetric quantization - along
-    with any combinations of these.
+    a given weight tensor `x` based on the specified parameters. It supports
+    grouped, per-channel, per-tensor, symmetric, and asymmetric quantization.
+
+    For grouped quantization (per_channel=True, group_size > 0), the output
+    shapes are [out_features, n_groups] where n_groups is the number of groups
+    along the in_features dimension.
 
     Args:
-        x: KerasTensor. The input tensor to quantize.
+        x: KerasTensor. The weight tensor to quantize with shape
+            [out_features, in_features].
         bits: int. The number of bits to quantize to (e.g., 4).
         symmetric: bool. Whether to use symmetric quantization.
         per_channel: bool. Whether to quantize per channel.
-        group_size: int. The group size for quantization.
-        weight: bool. Whether the input tensor is a weight tensor.
+        group_size: int. The group size for quantization. -1 means no grouping.
+        compute_dtype: str. The dtype for computation. Defaults to "float32".
 
     Returns:
         scale: KerasTensor. The scale tensor for quantization.
         zero: KerasTensor. The zero tensor for quantization.
         maxq: scalar. The maximum quantization value.
     """
+    # Input validation
     if x is None:
         raise ValueError(f"Input tensor {x} cannot be None.")
-
-    # For weights, we typically expect at least a 2D tensor.
-    if weight and len(x.shape) < 2:
+    if len(x.shape) < 2:
         raise ValueError(
             f"Input weight tensor {x} must have a rank of at "
             f"least 2, but got rank {len(x.shape)}."
         )
-
     if ops.size(x) == 0:
         raise ValueError("Input tensor 'x' cannot be empty.")
 
-    original_shape = x.shape
-
-    if per_channel:
-        if weight:
-            if group_size != -1:
-                input_reshaped = ops.reshape(x, [-1, group_size])
-            else:
-                input_reshaped = ops.reshape(x, [original_shape[0], -1])
-    else:  # per-tensor
-        input_reshaped = ops.reshape(x, [1, -1])
+    out_features, in_features = x.shape[0], x.shape[1]
 
-    # Find min/max values
-    min_values = ops.min(input_reshaped, axis=1)
-    max_values = ops.max(input_reshaped, axis=1)
+    # Determine number of groups for quantization
+    if per_channel and group_size > 0:
+        n_groups = (in_features + group_size - 1) // group_size
+    else:
+        n_groups = 1
+
+    # Compute min/max values based on quantization mode
+    if n_groups > 1:
+        # Grouped quantization: output shape [out_features, n_groups]
+        remainder = in_features % group_size
+        if remainder != 0:
+            pad_size = group_size - remainder
+            x = ops.pad(x, [[0, 0], [0, pad_size]], constant_values=0.0)
+
+        x_grouped = ops.reshape(x, [out_features, n_groups, group_size])
+        min_values = ops.min(x_grouped, axis=2)
+        max_values = ops.max(x_grouped, axis=2)
+    else:
+        # Per-channel or per-tensor: compute stats along rows
+        reduction_shape = [out_features, -1] if per_channel else [1, -1]
+        x_reshaped = ops.reshape(x, reduction_shape)
+        min_values = ops.min(x_reshaped, axis=1)
+        max_values = ops.max(x_reshaped, axis=1)
 
-    # Apply symmetric quantization logic if enabled
+    # Symmetric quantization: make range symmetric around zero
     if symmetric:
-        max_values = ops.maximum(ops.abs(min_values), max_values)
+        max_abs = ops.maximum(ops.abs(min_values), max_values)
         min_values = ops.where(
-            ops.less(min_values, 0), ops.negative(max_values), min_values
+            ops.less(min_values, 0), ops.negative(max_abs), min_values
         )
+        max_values = max_abs
 
-    # Ensure range is not zero to avoid division errors
+    # Ensure non-zero range to avoid division errors
     zero_range = ops.equal(min_values, max_values)
     min_values = ops.where(zero_range, ops.subtract(min_values, 1), min_values)
     max_values = ops.where(zero_range, ops.add(max_values, 1), max_values)
 
+    # Compute scale and zero-point
     maxq = ops.cast(ops.subtract(ops.power(2, bits), 1), compute_dtype)
-
-    # Calculate scale and zero-point
     scale = ops.divide(ops.subtract(max_values, min_values), maxq)
+    scale = ops.where(ops.less_equal(scale, 0), 1e-8, scale)
+
     if symmetric:
         zero = ops.full_like(scale, ops.divide(ops.add(maxq, 1), 2))
     else:
         zero = ops.round(ops.divide(ops.negative(min_values), scale))
 
-    # Ensure scale is non-zero
-    scale = ops.where(ops.less_equal(scale, 0), 1e-8, scale)
-
-    if weight:
-        # Per-channel, non-grouped case: simple reshape is correct.
-        if per_channel and group_size == -1:
-            scale = ops.reshape(scale, [-1, 1])
-            zero = ops.reshape(zero, [-1, 1])
-        elif not per_channel:
-            num_rows = original_shape[0]
-            scale = ops.tile(ops.reshape(scale, (1, 1)), (num_rows, 1))
-            zero = ops.tile(ops.reshape(zero, (1, 1)), (num_rows, 1))
-    if per_channel:
+    # Reshape output to [out_features, n_groups] or [out_features, 1]
+    if n_groups > 1:
+        pass  # Already [out_features, n_groups]
+    elif per_channel:
         scale = ops.reshape(scale, [-1, 1])
         zero = ops.reshape(zero, [-1, 1])
+    else:
+        # Per-tensor: tile single value to [out_features, 1]
+        scale = ops.tile(ops.reshape(scale, (1, 1)), (out_features, 1))
+        zero = ops.tile(ops.reshape(zero, (1, 1)), (out_features, 1))
 
-    zero = ops.cast(zero, "uint8")
-
-    return scale, zero, maxq
+    return scale, ops.cast(zero, "uint8"), maxq
 
 
 def quantize_with_zero_point(input_tensor, scale, zero, maxq):

keras/src/quantizers/utils.py

@@ -0,0 +1,23 @@
+import re
+
+
+def should_quantize_layer(layer, filters):
+    """Determines if a layer should be quantized based on filters.
+
+    Args:
+        layer: The layer to check.
+        filters: A regex string, a list of regex strings, or a callable.
+            If None, returns True.
+
+    Returns:
+        True if the layer should be quantized, False otherwise.
+    """
+    if filters is None:
+        return True
+    if isinstance(filters, str):
+        return bool(re.search(filters, layer.name))
+    if isinstance(filters, (list, tuple)):
+        return any(re.search(pat, layer.name) for pat in filters)
+    if callable(filters):
+        return filters(layer)
+    return True

keras/src/random/seed_generator.py

@@ -8,6 +8,8 @@ from keras.src.backend.common import global_state
 from keras.src.utils import jax_utils
 from keras.src.utils.naming import auto_name
 
+GLOBAL_SEED_GENERATOR = "global_seed_generator"
+
 
 @keras_export("keras.random.SeedGenerator")
 class SeedGenerator:
@@ -27,7 +29,7 @@ class SeedGenerator:
     a local `StateGenerator` with either a deterministic or random initial
     state.
 
-    Remark concerning the JAX backen: Note that the use of a local
+    Remark concerning the JAX backend: Note that the use of a local
     `StateGenerator` as seed argument is required for JIT compilation of
     RNG with the JAX backend, because the use of global state is not
     supported.
@@ -109,7 +111,7 @@ class SeedGenerator:
         return new_seed_value
 
     def get_config(self):
-        return {"seed": self._initial_seed}
+        return {"seed": self._initial_seed, "name": self.name}
 
     @classmethod
     def from_config(cls, config):
@@ -133,10 +135,10 @@ def global_seed_generator():
             "out = keras.random.normal(shape=(1,), seed=self.seed_generator)\n"
             "```"
         )
-    gen = global_state.get_global_attribute("global_seed_generator")
+    gen = global_state.get_global_attribute(GLOBAL_SEED_GENERATOR)
     if gen is None:
         gen = SeedGenerator()
-        global_state.set_global_attribute("global_seed_generator", gen)
+        global_state.set_global_attribute(GLOBAL_SEED_GENERATOR, gen)
     return gen