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

keras/src/layers/regularization/dropout.py

@@ -48,13 +48,55 @@ class Dropout(Layer):
             )
         self.rate = rate
         self.seed = seed
-        self.noise_shape = noise_shape
+        self.noise_shape = self._validate_noise_shape(noise_shape)
         if rate > 0:
             self.seed_generator = backend.random.SeedGenerator(seed)
         self.supports_masking = True
 
         self._build_at_init()
 
+    def _validate_noise_shape(self, noise_shape):
+        if noise_shape is None:
+            return None
+
+        if isinstance(noise_shape, str):
+            raise ValueError(
+                f"Invalid value received for argument `noise_shape`. "
+                f"Expected a tuple or list of integers. "
+                f"Received: noise_shape={noise_shape}"
+            )
+
+        if not isinstance(noise_shape, tuple):
+            try:
+                noise_shape = tuple(noise_shape)
+            except TypeError:
+                raise ValueError(
+                    f"Invalid value received for argument `noise_shape`. "
+                    f"Expected an iterable of integers "
+                    f"(e.g., a tuple or list). "
+                    f"Received: noise_shape={noise_shape}"
+                )
+
+        for i, dim in enumerate(noise_shape):
+            if dim is not None:
+                if not isinstance(dim, int):
+                    raise ValueError(
+                        f"Invalid value received for argument `noise_shape`. "
+                        f"Expected all elements to be integers or None. "
+                        f"Received element at index {i}: {dim} "
+                        f"(type: {type(dim).__name__})"
+                    )
+
+                if dim <= 0:
+                    raise ValueError(
+                        f"Invalid value received for argument `noise_shape`. "
+                        f"Expected all dimensions to be positive integers "
+                        f"or None. "
+                        f"Received negative or zero value at index {i}: {dim}"
+                    )
+
+        return noise_shape
+
     def call(self, inputs, training=False):
         if training and self.rate > 0:
             return backend.random.dropout(

keras/src/layers/rnn/rnn.py

@@ -212,6 +212,7 @@ class RNN(Layer):
         self.supports_masking = True
         self.input_spec = None
         self.states = None
+        self._expected_batch_size = None
 
         state_size = getattr(self.cell, "state_size", None)
         if state_size is None:
@@ -283,6 +284,9 @@ class RNN(Layer):
                         f"batch size: sequence.shape={sequences_shape}"
                     )
                 self._create_state_variables(sequences_shape[0])
+                self._expected_batch_size = ops.shape(
+                    tree.flatten(self.states)[0]
+                )[0]
 
     @tracking.no_automatic_dependency_tracking
     def _create_state_variables(self, batch_size):
@@ -382,6 +386,21 @@ class RNN(Layer):
                 initial_state = self.get_initial_state(
                     batch_size=ops.shape(sequences)[0]
                 )
+        if self.stateful:
+            actual_batch_size = sequences.shape[0]
+            if (
+                self._expected_batch_size is not None
+                and actual_batch_size is not None
+                and actual_batch_size != self._expected_batch_size
+            ):
+                raise ValueError(
+                    f"If an RNN is stateful, the batch size of the "
+                    f"input sequences must be the same as the batch "
+                    f"size of the initial state. \n"
+                    f"- Expected batch size: {self._expected_batch_size}\n"
+                    f"- Received batch size: {actual_batch_size}"
+                )
+
         # RNN expect the states in a list, even if single state.
         if not tree.is_nested(initial_state):
             initial_state = [initial_state]

keras/src/losses/loss.py

@@ -211,7 +211,7 @@ def apply_mask(sample_weight, mask, dtype, reduction):
                 dtype,
             )
             valid = ops.sum(mask)  # May be 0!
-            mask *= total / (valid + backend.epsilon())
+            mask *= ops.divide_no_nan(total, valid)
 
         if sample_weight is not None:
             sample_weight = ops.cast(sample_weight, dtype=dtype)

keras/src/losses/losses.py

@@ -73,6 +73,14 @@ class MeanSquaredError(LossFunctionWrapper):
             `"float32"` unless set to different value
             (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
             provided, then the `compute_dtype` will be utilized.
+
+    Examples:
+
+    >>> y_true = keras.ops.array([1.0, 0.0, 1.0])
+    >>> y_pred = keras.ops.array([0.9, 0.1, 0.8])
+    >>> loss = keras.losses.MeanSquaredError()
+    >>> loss(y_true, y_pred)
+    0.02
     """
 
     def __init__(
@@ -114,6 +122,14 @@ class MeanAbsoluteError(LossFunctionWrapper):
             `"float32"` unless set to different value
             (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
             provided, then the `compute_dtype` will be utilized.
+
+    Examples:
+
+    >>> y_true = keras.ops.array([1.0, 0.3, 1.0])
+    >>> y_pred = keras.ops.array([1.9, 0.3, 1.8])
+    >>> loss = keras.losses.MeanAbsoluteError()
+    >>> loss(y_true, y_pred)
+    0.5666667
     """
 
     def __init__(
@@ -155,6 +171,14 @@ class MeanAbsolutePercentageError(LossFunctionWrapper):
             `"float32"` unless set to different value
             (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
             provided, then the `compute_dtype` will be utilized.
+
+    Examples:
+
+    >>> y_true = keras.ops.array([100.0, 200.0, 300.0])
+    >>> y_pred = keras.ops.array([90.0, 210.0, 310.0])
+    >>> loss = keras.losses.MeanAbsolutePercentageError()
+    >>> loss(y_true, y_pred)
+    6.111111
     """
 
     def __init__(

keras/src/metrics/confusion_metrics.py

@@ -654,7 +654,7 @@ class SensitivitySpecificityBase(Metric):
         Args:
             constrained: Over these values the constraint is specified. A rank-1
                 tensor.
-            dependent: From these values the maximum that satiesfies the
+            dependent: From these values the maximum that satisfies the
                 constraint is selected. Values in this tensor and in
                 `constrained` are linked by having the same threshold at each
                 position, hence this tensor must have the same shape.
@@ -664,11 +664,12 @@ class SensitivitySpecificityBase(Metric):
         Returns:
             maximal dependent value, if no value satisfies the constraint 0.0.
         """
-        feasible = ops.nonzero(predicate(constrained, self.value))
-        feasible_exists = ops.greater(ops.size(feasible), 0)
-        max_dependent = ops.max(ops.take(dependent, feasible), initial=0)
-
-        return ops.where(feasible_exists, max_dependent, 0.0)
+        feasible = predicate(constrained, self.value)
+        # Mask values based on whether they satisfy the constraint and take max.
+        return ops.max(
+            ops.multiply(dependent, ops.cast(feasible, dependent.dtype)),
+            initial=0,
+        )
 
 
 @keras_export("keras.metrics.SensitivityAtSpecificity")

keras/src/models/cloning.py

@@ -293,10 +293,12 @@ def _clone_sequential_model(model, clone_function, input_tensors=None):
         input_name = ref_input_layer.name
         input_batch_shape = ref_input_layer.batch_shape
         input_dtype = ref_input_layer._dtype
+        input_optional = ref_input_layer.optional
     else:
         input_name = None
         input_dtype = None
         input_batch_shape = None
+        input_optional = False
 
     if input_tensors is not None:
         if isinstance(input_tensors, (list, tuple)):
@@ -313,6 +315,7 @@ def _clone_sequential_model(model, clone_function, input_tensors=None):
         inputs = Input(
             tensor=input_tensors,
             name=input_name,
+            optional=input_optional,
         )
         new_layers = [inputs] + new_layers
     else:
@@ -321,6 +324,7 @@ def _clone_sequential_model(model, clone_function, input_tensors=None):
                 batch_shape=input_batch_shape,
                 dtype=input_dtype,
                 name=input_name,
+                optional=input_optional,
             )
             new_layers = [inputs] + new_layers
     cloned_model = Sequential(

keras/src/models/functional.py

@@ -254,9 +254,9 @@ class Functional(Function, Model):
         return converted
 
     def _adjust_input_rank(self, flat_inputs):
-        flat_ref_shapes = [x.shape for x in self._inputs]
         adjusted = []
-        for x, ref_shape in zip(flat_inputs, flat_ref_shapes):
+        for i, x in enumerate(flat_inputs):
+            ref_shape = self._inputs[i].shape
             if x is None:
                 adjusted.append(x)
                 continue
@@ -273,8 +273,11 @@ class Functional(Function, Model):
                 if ref_shape[-1] == 1:
                     adjusted.append(ops.expand_dims(x, axis=-1))
                     continue
+            flat_paths_and_inputs = tree.flatten_with_path(self._inputs_struct)
+            path = ".".join(str(p) for p in flat_paths_and_inputs[i][0])
             raise ValueError(
-                f"Invalid input shape for input {x}. Expected shape "
+                f"Invalid input shape for input {x} with name "
+                f"'{self._inputs[i].name}' and path '{path}'. Expected shape "
                 f"{ref_shape}, but input has incompatible shape {x.shape}"
             )
         # Add back metadata.
@@ -832,11 +835,16 @@ def clone_graph_nodes(inputs, outputs):
             kt_id_mapping[id(kt_input)] = kt_input
         else:
             # We need to create a new Keras tensor for any intermediate tensor
+            original_op = kt_input._keras_history.operation
+            optional = False
+            if isinstance(original_op, InputLayer):
+                optional = original_op.optional
             cloned_input = Input(
                 batch_shape=kt_input.shape,
                 dtype=kt_input.dtype,
                 sparse=kt_input.sparse,
                 name=f"{kt_input.name}CLONE",
+                optional=optional,
             )
             cloned_inputs.append(cloned_input)
             kt_id_mapping[id(kt_input)] = cloned_input

keras/src/models/model.py

@@ -2,14 +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
@@ -422,19 +424,99 @@ 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.
+
+        Example:
+
+        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:
@@ -443,27 +525,20 @@ class Model(Trainer, base_trainer.Trainer, Layer):
                 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}"
-            )
-
-        if mode == "gptq":
-            if not isinstance(config, GPTQConfig):
+        if filters is not None:
+            if not isinstance(filters, (str, Callable, list, tuple)):
                 raise ValueError(
-                    "Mode 'gptq' requires a valid `config` argument of type "
-                    f"`GPTQConfig`. Received: {type(config)}"
+                    "The `filters` argument must be a regex string, a list of "
+                    "regex strings, or a callable. Received: "
+                    f"{type(filters)}"
                 )
-        elif config is not None:
-            # All other modes must not receive a config
-            raise ValueError(
-                f"The `config` argument is only supported for 'gptq' mode, "
-                f"but received mode='{mode}' and a non-None config."
-            )
 
         graph_modified = False
         for layer in self._flatten_layers():
+            # 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, config=config)
@@ -473,8 +548,29 @@ class Model(Trainer, base_trainer.Trainer, Layer):
                 except AttributeError:
                     pass
 
-        if mode == "gptq":
-            gptq_quantize(self, config)
+        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:
@@ -569,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.
@@ -593,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.
@@ -627,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,
@@ -663,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):
@@ -863,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}")