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

keras/src/models/model.py

@@ -2,13 +2,16 @@ import inspect
 import json
 import typing
 import warnings
+from collections.abc import Callable
 
 from keras.src import backend
 from keras.src import utils
 from keras.src.api_export import keras_export
 from keras.src.layers.layer import Layer
 from keras.src.models.variable_mapping import map_saveable_variables
-from keras.src.quantizers.gptq_config import GPTQConfig
+from keras.src.quantizers.awq_core import awq_quantize
+from keras.src.quantizers.gptq_core import gptq_quantize
+from keras.src.quantizers.utils import should_quantize_layer
 from keras.src.saving import saving_api
 from keras.src.trainers import trainer as base_trainer
 from keras.src.utils import summary_utils
@@ -421,62 +424,168 @@ class Model(Trainer, base_trainer.Trainer, Layer):
             **kwargs,
         )
 
-    def quantize(self, mode, config=None, **kwargs):
+    def get_quantization_layer_structure(self, mode=None):
+        """Returns the quantization structure for the model.
+
+        This method is intended to be overridden by model authors to provide
+        topology information required for structure-aware quantization modes
+        like 'gptq'.
+
+        Args:
+            mode: The quantization mode.
+
+        Returns:
+            A dictionary describing the topology, e.g.:
+            `{'pre_block_layers': [list], 'sequential_blocks': [list]}`
+            or `None` if the mode does not require structure or is not
+            supported. `'pre_block_layers'` is a list of layers that
+            the inputs should be passed through, before being passed to
+            the sequential blocks. For example, inputs to an LLM must
+            first be passed through an embedding layer, followed by
+            the transformer.
+        """
+        del mode  # Unused.
+        return None
+
+    def quantize(self, mode=None, config=None, filters=None, **kwargs):
         """Quantize the weights of the model.
 
         Note that the model must be built first before calling this method.
-        `quantize` will recursively call `quantize(mode)` in all layers and
+        `quantize` will recursively call `quantize(...)` in all layers and
         will be skipped if the layer doesn't implement the function.
 
+        This method can be called by passing a `mode` string, which uses the
+        default configuration for that mode. Alternatively, a `config` object
+        can be passed to customize the behavior of the quantization (e.g. to
+        use specific quantizers for weights or activations).
+
         Args:
-            mode: The mode of the quantization. Only 'int8' is supported at this
-                time.
-        """
-        from keras.src.dtype_policies import QUANTIZATION_MODES
+            mode: The mode of the quantization. Supported modes are:
+                `"int8"`, `"int4"`, `"float8"`, `"gptq"`. This is
+                optional if `config` is provided.
+            config: The configuration object specifying additional
+                quantization options. This argument allows to configure
+                the weight and activation quantizers. be an instance of
+                `keras.quantizers.QuantizationConfig`.
+            filters: Optional filters to apply to the quantization. Can be a
+                regex string, a list of regex strings, or a callable. Only the
+                layers which match the filter conditions will be quantized.
+            **kwargs: Additional keyword arguments.
 
-        if mode == "gptq":
-            if not isinstance(config, GPTQConfig):
-                raise ValueError(
-                    "The `config` argument must be of type "
-                    "`keras.quantizers.GPTQConfig`."
-                )
-            # The config object's own quantize method drives the process
-            config.quantize(self)
-            return
+        Example:
 
-        # For all other modes, verify that a config object was not passed.
-        if config is not None:
-            raise ValueError(
-                f"The `config` argument is only supported for 'gptq' mode, "
-                f"but received mode='{mode}'."
-            )
+        Quantize a model to int8 with default configuration:
+
+        ```python
+        # Build the model
+        model = keras.Sequential([
+            keras.Input(shape=(10,)),
+            keras.layers.Dense(10),
+        ])
+        model.build((None, 10))
 
+        # Quantize with default int8 config
+        model.quantize("int8")
+        ```
+
+        Quantize a model to int8 with a custom configuration:
+
+        ```python
+        from keras.quantizers import Int8QuantizationConfig
+        from keras.quantizers import AbsMaxQuantizer
+
+        # Build the model
+        model = keras.Sequential([
+            keras.Input(shape=(10,)),
+            keras.layers.Dense(10),
+        ])
+        model.build((None, 10))
+
+        # Create a custom config
+        config = Int8QuantizationConfig(
+            weight_quantizer=AbsMaxQuantizer(
+                axis=0,
+                value_range=(-127, 127)
+            ),
+            activation_quantizer=AbsMaxQuantizer(
+                axis=-1,
+                value_range=(-127, 127)
+            ),
+        )
+
+        # Quantize with custom config
+        model.quantize(config=config)
+        ```
+        """
+        # Validate inputs.
         type_check = kwargs.pop("type_check", True)
         if kwargs:
             raise ValueError(
                 "Unrecognized keyword arguments "
                 f"passed to {self.__class__.__name__}: {kwargs}"
             )
-        if mode not in QUANTIZATION_MODES:
-            raise ValueError(
-                "Invalid quantization mode. "
-                f"Expected one of {QUANTIZATION_MODES}. Received: mode={mode}"
-            )
-        mode_changed = False
+
+        if filters is not None:
+            if not isinstance(filters, (str, Callable, list, tuple)):
+                raise ValueError(
+                    "The `filters` argument must be a regex string, a list of "
+                    "regex strings, or a callable. Received: "
+                    f"{type(filters)}"
+                )
+
+        graph_modified = False
         for layer in self._flatten_layers():
-            list_of_sublayers = list(layer._flatten_layers())
-            if len(list_of_sublayers) == 1:  # leaves of the model
+            # Apply filters
+            if not should_quantize_layer(layer, filters):
+                continue
+
+            if len(list(layer._flatten_layers())) == 1:
                 try:
-                    layer.quantize(mode, type_check=type_check)
-                    mode_changed = True
+                    layer.quantize(mode, type_check=type_check, config=config)
+                    graph_modified = True
                 except NotImplementedError as e:
                     warnings.warn(str(e))
-        # We need to set these functions to `None` to remake them for changed
-        # call function
-        if mode_changed:
+                except AttributeError:
+                    pass
+
+        if mode in ["gptq", "awq"]:
+            # Resolve model structure.
+            # 1. If quantization_layer_structure is provided inside the config,
+            # use that.
+            structure = config.quantization_layer_structure
+            # 2. If no layer structure is provided in the config, try to fetch
+            # it using the `get_quantization_layer_structure` hook.
+            if structure is None:
+                structure = self.get_quantization_layer_structure(mode)
+
+            if structure is None:
+                raise ValueError(
+                    f"For {mode=}, a valid quantization structure must be "
+                    "provided either via `config.quantization_layer_structure` "
+                    "or by overriding "
+                    "`model.get_quantization_layer_structure(mode)`. The "
+                    "structure should be a dictionary with keys "
+                    "'pre_block_layers' and 'sequential_blocks'."
+                )
+            if mode == "gptq":
+                gptq_quantize(config, structure, filters=filters)
+            elif mode == "awq":
+                awq_quantize(config, structure, filters=filters)
+
+        # If any layer was changed, we must rebuild the execution functions.
+        if graph_modified:
             self.train_function = None
             self.test_function = None
             self.predict_function = None
+            self._post_quantize(mode, **kwargs)
+
+    def _post_quantize(self, mode, **kwargs):
+        if backend.backend() == "torch":
+            # We need to manually retrack `torch_params`.
+            # The reason is that after quantization, the removed variables are
+            # still referenced by `torch_params` and cannot be gc.
+            for layer in self._flatten_layers():
+                layer._track_variables()
 
     def build_from_config(self, config):
         if not config:
@@ -556,8 +665,8 @@ class Model(Trainer, base_trainer.Trainer, Layer):
             filepath: `str` or `pathlib.Path` object. The path to save the
                 artifact.
             format: `str`. The export format. Supported values:
-                `"tf_saved_model"` and `"onnx"`.  Defaults to
-                `"tf_saved_model"`.
+                `"tf_saved_model"`, `"onnx"`, `"openvino"`, and `"litert"`.
+                Defaults to `"tf_saved_model"`.
             verbose: `bool`. Whether to print a message during export. Defaults
                 to `None`, which uses the default value set by different
                 backends and formats.
@@ -580,6 +689,13 @@ class Model(Trainer, base_trainer.Trainer, Layer):
                     provided, they will be automatically computed.
                 - `opset_version`: Optional `int`. Specific to `format="onnx"`.
                     An integer value that specifies the ONNX opset version.
+                - LiteRT-specific options: Optional keyword arguments specific
+                    to `format="litert"`. These are passed directly to the
+                    TensorFlow Lite converter and include options like
+                    `optimizations`, `representative_dataset`,
+                    `experimental_new_quantizer`, `allow_custom_ops`,
+                    `enable_select_tf_ops`, etc. See TensorFlow Lite
+                    documentation for all available options.
 
         **Note:** This feature is currently supported only with TensorFlow, JAX
         and Torch backends.
@@ -614,18 +730,41 @@ class Model(Trainer, base_trainer.Trainer, Layer):
         }
         predictions = ort_session.run(None, ort_inputs)
         ```
+
+        Here's how to export a LiteRT (TFLite) for inference.
+
+        ```python
+        # Export the model as a LiteRT artifact
+        model.export("path/to/location", format="litert")
+
+        # Load the artifact in a different process/environment
+        interpreter = tf.lite.Interpreter(model_path="path/to/location")
+        interpreter.allocate_tensors()
+        interpreter.set_tensor(
+            interpreter.get_input_details()[0]['index'], input_data
+        )
+        interpreter.invoke()
+        output_data = interpreter.get_tensor(
+            interpreter.get_output_details()[0]['index']
+        )
+        ```
         """
+        from keras.src.export import export_litert
         from keras.src.export import export_onnx
         from keras.src.export import export_openvino
         from keras.src.export import export_saved_model
 
-        available_formats = ("tf_saved_model", "onnx", "openvino")
+        available_formats = ("tf_saved_model", "onnx", "openvino", "litert")
         if format not in available_formats:
             raise ValueError(
                 f"Unrecognized format={format}. Supported formats are: "
                 f"{list(available_formats)}."
             )
 
+        # Check if LiteRT export is available (requires TensorFlow backend)
+        if format == "litert" and backend.backend() != "tensorflow":
+            raise ImportError("LiteRT export requires TensorFlow backend.")
+
         if format == "tf_saved_model":
             export_saved_model(
                 self,
@@ -650,6 +789,13 @@ class Model(Trainer, base_trainer.Trainer, Layer):
                 input_signature=input_signature,
                 **kwargs,
             )
+        elif format == "litert":
+            export_litert(
+                self,
+                filepath,
+                input_signature=input_signature,
+                **kwargs,
+            )
 
     @classmethod
     def from_config(cls, config, custom_objects=None):
@@ -850,13 +996,18 @@ class Model(Trainer, base_trainer.Trainer, Layer):
                     self.non_trainable_variables, path_value_dict
                 )
             elif k == "optimizer_variables":
-                self._assign_variable_values(
-                    self.optimizer.variables, path_value_dict
-                )
+                if hasattr(self, "optimizer") and self.optimizer is not None:
+                    self._assign_variable_values(
+                        self.optimizer.variables, path_value_dict
+                    )
             elif k == "metrics_variables":
-                self._assign_variable_values(
-                    self.metrics_variables, path_value_dict
-                )
+                if (
+                    hasattr(self, "metrics_variables")
+                    and self.metrics_variables
+                ):
+                    self._assign_variable_values(
+                        self.metrics_variables, path_value_dict
+                    )
             else:
                 raise ValueError(f"Unknown variable name: {k}")
 

keras/src/ops/image.py

@@ -565,6 +565,8 @@ class ExtractPatches(Operation):
         if isinstance(size, int):
             size = (size, size)
         self.size = size
+        if strides is None:
+            strides = size
         self.strides = strides
         self.dilation_rate = dilation_rate
         self.padding = padding
@@ -583,8 +585,6 @@ class ExtractPatches(Operation):
     def compute_output_spec(self, images):
         images_shape = list(images.shape)
         original_ndim = len(images_shape)
-        if not self.strides:
-            strides = (self.size[0], self.size[1])
         if self.data_format == "channels_last":
             channels_in = images_shape[-1]
         else:
@@ -597,7 +597,7 @@ class ExtractPatches(Operation):
             images_shape,
             filters,
             kernel_size,
-            strides=strides,
+            strides=self.strides,
             padding=self.padding,
             data_format=self.data_format,
             dilation_rate=self.dilation_rate,
@@ -616,42 +616,98 @@ def extract_patches(
     padding="valid",
     data_format=None,
 ):
-    """Extracts patches from the image(s).
+    """Extracts patches from the image(s) or volume(s).
+
+    This function supports both 2D and 3D patch extraction based on the
+    `size` argument length, similar to how `keras.ops.conv` handles
+    different dimensions.
 
     Args:
-        images: Input image or batch of images. Must be 3D or 4D.
-        size: Patch size int or tuple (patch_height, patch_width)
-        strides: strides along height and width. If not specified, or
-            if `None`, it defaults to the same value as `size`.
-        dilation_rate: This is the input stride, specifying how far two
-            consecutive patch samples are in the input. For value other than 1,
-            strides must be 1. NOTE: `strides > 1` is not supported in
-            conjunction with `dilation_rate > 1`
+        images: Input image/volume or batch of images/volumes.
+            For 2D patches: 3D `(H, W, C)` or 4D `(N, H, W, C)`.
+            For 3D patches: 4D `(D, H, W, C)` or 5D `(N, D, H, W, C)`.
+        size: Patch size as int or tuple.
+            Length 2 tuple `(patch_height, patch_width)` or int for 2D patches.
+            Length 3 tuple `(patch_depth, patch_height, patch_width)` for
+            3D patches.
+        strides: Strides for patch extraction. If not specified, defaults
+            to `size` (non-overlapping patches).
+        dilation_rate: Dilation rate for patch extraction. Note that
+            `dilation_rate > 1` is not supported with `strides > 1`.
         padding: The type of padding algorithm to use: `"same"` or `"valid"`.
         data_format: A string specifying the data format of the input tensor.
             It can be either `"channels_last"` or `"channels_first"`.
-            `"channels_last"` corresponds to inputs with shape
-            `(batch, height, width, channels)`, while `"channels_first"`
-            corresponds to inputs with shape `(batch, channels, height, width)`.
-            If not specified, the value will default to
-            `keras.config.image_data_format`.
+            If not specified, defaults to `keras.config.image_data_format`.
 
     Returns:
-        Extracted patches 3D (if not batched) or 4D (if batched)
+        Extracted patches with shape depending on input and `size`:
+        - 2D patches: 3D (unbatched) or 4D (batched)
+        - 3D patches: 4D (unbatched) or 5D (batched)
 
     Examples:
 
+    >>> # 2D patches from batch of images
     >>> image = np.random.random(
     ...     (2, 20, 20, 3)
-    ... ).astype("float32") # batch of 2 RGB images
+    ... ).astype("float32")
     >>> patches = keras.ops.image.extract_patches(image, (5, 5))
     >>> patches.shape
     (2, 4, 4, 75)
-    >>> image = np.random.random((20, 20, 3)).astype("float32") # 1 RGB image
+
+    >>> # 2D patches from single image
+    >>> image = np.random.random((20, 20, 3)).astype("float32")
     >>> patches = keras.ops.image.extract_patches(image, (3, 3), (1, 1))
     >>> patches.shape
     (18, 18, 27)
+
+    >>> # 3D patches from batch of volumes
+    >>> volumes = np.random.random(
+    ...     (2, 10, 10, 10, 3)
+    ... ).astype("float32")
+    >>> patches = keras.ops.image.extract_patches(volumes, (3, 3, 3))
+    >>> patches.shape
+    (2, 3, 3, 3, 81)
+
+    >>> # 3D patches from single volume
+    >>> volume = np.random.random((10, 10, 10, 3)).astype("float32")
+    >>> patches = keras.ops.image.extract_patches(volume, (3, 3, 3))
+    >>> patches.shape
+    (3, 3, 3, 81)
     """
+    # Validate size argument
+    if not isinstance(size, int):
+        if not isinstance(size, (tuple, list)):
+            raise TypeError(
+                "Invalid `size` argument. Expected an int or a tuple. "
+                f"Received: size={size} of type {type(size).__name__}"
+            )
+        if len(size) not in (2, 3):
+            raise ValueError(
+                "Invalid `size` argument. Expected a tuple of length 2 or 3. "
+                f"Received: size={size} with length {len(size)}"
+            )
+
+    # Determine 2D vs 3D based on size argument
+    if not isinstance(size, int) and len(size) == 3:
+        # 3D patch extraction
+        if any_symbolic_tensors((images,)):
+            return ExtractPatches3D(
+                size=size,
+                strides=strides,
+                dilation_rate=dilation_rate,
+                padding=padding,
+                data_format=data_format,
+            ).symbolic_call(images)
+        return _extract_patches_3d(
+            images,
+            size,
+            strides,
+            dilation_rate,
+            padding,
+            data_format=data_format,
+        )
+
+    # 2D patch extraction (default)
     if any_symbolic_tensors((images,)):
         return ExtractPatches(
             size=size,
@@ -712,6 +768,187 @@ def _extract_patches(
     return patches
 
 
+class ExtractPatches3D(Operation):
+    def __init__(
+        self,
+        size,
+        strides=None,
+        dilation_rate=1,
+        padding="valid",
+        data_format=None,
+        *,
+        name=None,
+    ):
+        super().__init__(name=name)
+        if isinstance(size, int):
+            size = (size, size, size)
+        elif len(size) != 3:
+            raise TypeError(
+                "Invalid `size` argument. Expected an "
+                f"int or a tuple of length 3. Received: size={size}"
+            )
+        self.size = size
+        if strides is not None:
+            if isinstance(strides, int):
+                strides = (strides, strides, strides)
+            elif len(strides) != 3:
+                raise ValueError(f"Invalid `strides` argument. Got: {strides}")
+        else:
+            strides = size
+        self.strides = strides
+        self.dilation_rate = dilation_rate
+        self.padding = padding
+        self.data_format = backend.standardize_data_format(data_format)
+
+    def call(self, volumes):
+        return _extract_patches_3d(
+            volumes,
+            self.size,
+            self.strides,
+            self.dilation_rate,
+            self.padding,
+            self.data_format,
+        )
+
+    def compute_output_spec(self, volumes):
+        volumes_shape = list(volumes.shape)
+        original_ndim = len(volumes_shape)
+        strides = self.strides
+        if self.data_format == "channels_last":
+            channels_in = volumes_shape[-1]
+        else:
+            channels_in = volumes_shape[-4]
+        if original_ndim == 4:
+            volumes_shape = [1] + volumes_shape
+        filters = self.size[0] * self.size[1] * self.size[2] * channels_in
+        kernel_size = (self.size[0], self.size[1], self.size[2])
+        out_shape = compute_conv_output_shape(
+            volumes_shape,
+            filters,
+            kernel_size,
+            strides=strides,
+            padding=self.padding,
+            data_format=self.data_format,
+            dilation_rate=self.dilation_rate,
+        )
+        if original_ndim == 4:
+            out_shape = out_shape[1:]
+        return KerasTensor(shape=out_shape, dtype=volumes.dtype)
+
+
+def _extract_patches_3d(
+    volumes,
+    size,
+    strides=None,
+    dilation_rate=1,
+    padding="valid",
+    data_format=None,
+):
+    if isinstance(size, int):
+        patch_d = patch_h = patch_w = size
+    elif len(size) == 3:
+        patch_d, patch_h, patch_w = size
+    else:
+        raise TypeError(
+            "Invalid `size` argument. Expected an "
+            f"int or a tuple of length 3. Received: size={size}"
+        )
+    if strides is None:
+        strides = size
+    if isinstance(strides, int):
+        strides = (strides, strides, strides)
+    if len(strides) != 3:
+        raise ValueError(f"Invalid `strides` argument. Got: {strides}")
+    data_format = backend.standardize_data_format(data_format)
+    if data_format == "channels_last":
+        channels_in = volumes.shape[-1]
+    elif data_format == "channels_first":
+        channels_in = volumes.shape[-4]
+    out_dim = patch_d * patch_w * patch_h * channels_in
+    kernel = backend.numpy.eye(out_dim, dtype=volumes.dtype)
+    kernel = backend.numpy.reshape(
+        kernel, (patch_d, patch_h, patch_w, channels_in, out_dim)
+    )
+    _unbatched = False
+    if len(volumes.shape) == 4:
+        _unbatched = True
+        volumes = backend.numpy.expand_dims(volumes, axis=0)
+    patches = backend.nn.conv(
+        inputs=volumes,
+        kernel=kernel,
+        strides=strides,
+        padding=padding,
+        data_format=data_format,
+        dilation_rate=dilation_rate,
+    )
+    if _unbatched:
+        patches = backend.numpy.squeeze(patches, axis=0)
+    return patches
+
+
+@keras_export("keras.ops.image.extract_patches_3d")
+def extract_patches_3d(
+    volumes,
+    size,
+    strides=None,
+    dilation_rate=1,
+    padding="valid",
+    data_format=None,
+):
+    """Extracts patches from the volume(s).
+
+    Args:
+        volumes: Input volume or batch of volumes. Must be 4D or 5D.
+        size: Patch size int or tuple (patch_depth, patch_height, patch_width)
+        strides: strides along depth, height, and width. If not specified, or
+            if `None`, it defaults to the same value as `size`.
+        dilation_rate: This is the input stride, specifying how far two
+            consecutive patch samples are in the input. Note that using
+            `dilation_rate > 1` is not supported in conjunction with
+            `strides > 1` on the TensorFlow backend.
+        padding: The type of padding algorithm to use: `"same"` or `"valid"`.
+        data_format: A string specifying the data format of the input tensor.
+            It can be either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, depth, height, width, channels)`, while `"channels_first"`
+            corresponds to inputs with shape
+            `(batch, channels, depth, height, width)`. If not specified,
+             the value will default to `keras.config.image_data_format()`.
+
+    Returns:
+        Extracted patches 4D (if not batched) or 5D (if batched)
+
+    Examples:
+
+    >>> import numpy as np
+    >>> import keras
+    >>> # Batched case
+    >>> volumes = np.random.random(
+    ...     (2, 10, 10, 10, 3)
+    ... ).astype("float32") # batch of 2 volumes
+    >>> patches = keras.ops.image.extract_patches_3d(volumes, (3, 3, 3))
+    >>> patches.shape
+    (2, 3, 3, 3, 81)
+    >>> # Unbatched case
+    >>> volume = np.random.random((10, 10, 10, 3)).astype("float32") # 1 volume
+    >>> patches = keras.ops.image.extract_patches_3d(volume, (3, 3, 3))
+    >>> patches.shape
+    (3, 3, 3, 81)
+    """
+    if any_symbolic_tensors((volumes,)):
+        return ExtractPatches3D(
+            size=size,
+            strides=strides,
+            dilation_rate=dilation_rate,
+            padding=padding,
+            data_format=data_format,
+        ).symbolic_call(volumes)
+
+    return _extract_patches_3d(
+        volumes, size, strides, dilation_rate, padding, data_format=data_format
+    )
+
+
 class MapCoordinates(Operation):
     def __init__(self, order, fill_mode="constant", fill_value=0, *, name=None):
         super().__init__(name=name)