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

keras/src/layers/pooling/adaptive_average_pooling1d.py

@@ -0,0 +1,65 @@
+"""Adaptive Average Pooling 1D layer."""
+
+from keras.src.api_export import keras_export
+from keras.src.layers.pooling.base_adaptive_pooling import (
+    BaseAdaptiveAveragePooling,
+)
+
+
+@keras_export("keras.layers.AdaptiveAveragePooling1D")
+class AdaptiveAveragePooling1D(BaseAdaptiveAveragePooling):
+    """Adaptive average pooling operation for 1D temporal or spatial data.
+
+    This layer applies an adaptive average pooling operation, which pools the
+    input such that the output has a target length specified by `output_size`,
+    regardless of the input length. The kernel size and stride are automatically
+    computed to achieve the target output size.
+
+    Args:
+        output_size: Integer specifying the target output length.
+        data_format: string, either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, length, channels)`.
+            `"channels_first"` corresponds to inputs with shape
+            `(batch, channels, length)`.
+            Defaults to the value found in your Keras config file at
+            `~/.keras/keras.json`. If never set, `"channels_last"` is used.
+
+    Input shape:
+        - If `data_format="channels_last"`: 3D tensor
+            `(batch_size, length, channels)`
+        - If `data_format="channels_first"`: 3D tensor
+            `(batch_size, channels, length)`
+
+    Output shape:
+        - If `data_format="channels_last"`:
+            `(batch_size, output_length, channels)`
+        - If `data_format="channels_first"`:
+            `(batch_size, channels, output_length)`
+
+    Examples:
+        >>> import numpy as np
+        >>> input_seq = np.random.rand(1, 64, 3)
+        >>> layer = AdaptiveAveragePooling1D(output_size=32)
+        >>> output_seq = layer(input_seq)
+        >>> output_seq.shape
+        (1, 32, 3)
+    """
+
+    def __init__(self, output_size, data_format=None, **kwargs):
+        if isinstance(output_size, int):
+            output_size = (output_size,)
+        elif isinstance(output_size, (tuple, list)):
+            if len(output_size) != 1:
+                raise ValueError(
+                    f"For 1D input, `output_size` tuple must have length 1. "
+                    f"Received: {output_size}"
+                )
+            output_size = tuple(output_size)
+        else:
+            raise TypeError(
+                f"`output_size` must be an integer or tuple of 1 integer. "
+                f"Received: {output_size} of type {type(output_size)}"
+            )
+
+        super().__init__(output_size, data_format, **kwargs)

keras/src/layers/pooling/adaptive_average_pooling2d.py

@@ -0,0 +1,62 @@
+"""Adaptive Average Pooling 2D layer."""
+
+from keras.src.api_export import keras_export
+from keras.src.layers.pooling.base_adaptive_pooling import (
+    BaseAdaptiveAveragePooling,
+)
+
+
+@keras_export("keras.layers.AdaptiveAveragePooling2D")
+class AdaptiveAveragePooling2D(BaseAdaptiveAveragePooling):
+    """Adaptive average pooling operation for 2D spatial data.
+
+    This layer applies an adaptive average pooling operation, which pools the
+    input such that the output has a target spatial size specified by
+    `output_size`, regardless of the input spatial size. The kernel size
+    and stride are automatically computed to achieve the target output size.
+
+    Args:
+        output_size: Integer or tuple of 2 integers specifying the
+            target output size.
+            If an integer, the same value is used for both height and width.
+        data_format: string, either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, height, width, channels)`.
+            `"channels_first"` corresponds to inputs with shape
+            `(batch, channels, height, width)`.
+            Defaults to the value found in your Keras config file at
+            `~/.keras/keras.json`. If never set, `"channels_last"` is used.
+
+    Input shape:
+        - If `data_format="channels_last"`: 4D tensor
+            `(batch_size, height, width, channels)`
+        - If `data_format="channels_first"`: 4D tensor
+            `(batch_size, channels, height, width)`
+
+    Output shape:
+        - If `data_format="channels_last"`:
+            `(batch_size, output_height, output_width, channels)`
+        - If `data_format="channels_first"`:
+            `(batch_size, channels, output_height, output_width)`
+
+    Examples:
+        >>> import numpy as np
+        >>> input_img = np.random.rand(1, 64, 64, 3)
+        >>> layer = AdaptiveAveragePooling2D(output_size=32)
+        >>> output_img = layer(input_img)
+        >>> output_img.shape
+        (1, 32, 32, 3)
+    """
+
+    def __init__(self, output_size, data_format=None, **kwargs):
+        if isinstance(output_size, int):
+            output_size_tuple = (output_size, output_size)
+        elif isinstance(output_size, (tuple, list)) and len(output_size) == 2:
+            output_size_tuple = tuple(output_size)
+        else:
+            raise TypeError(
+                f"`output_size` must be an integer or (height, width) tuple. "
+                f"Received: {output_size} of type {type(output_size)}"
+            )
+
+        super().__init__(output_size_tuple, data_format, **kwargs)

keras/src/layers/pooling/adaptive_average_pooling3d.py

@@ -0,0 +1,63 @@
+"""Adaptive Average Pooling 3D layer."""
+
+from keras.src.api_export import keras_export
+from keras.src.layers.pooling.base_adaptive_pooling import (
+    BaseAdaptiveAveragePooling,
+)
+
+
+@keras_export("keras.layers.AdaptiveAveragePooling3D")
+class AdaptiveAveragePooling3D(BaseAdaptiveAveragePooling):
+    """Adaptive average pooling operation for 3D volumetric data.
+
+    This layer applies an adaptive average pooling operation, which pools the
+    input such that the output has a target spatial size specified by
+    `output_size`, regardless of the input spatial size. The kernel size
+    and stride are automatically computed to achieve the target output size.
+
+    Args:
+        output_size: Integer or tuple of 3 integers specifying the
+            target output size.
+            If an integer, the same value is used for depth, height, and width.
+        data_format: string, either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, depth, height, width, channels)`.
+            `"channels_first"` corresponds to inputs with shape
+            `(batch, channels, depth, height, width)`.
+            Defaults to the value found in your Keras config file at
+            `~/.keras/keras.json`. If never set, `"channels_last"` is used.
+
+    Input shape:
+        - If `data_format="channels_last"`: 5D tensor
+            `(batch_size, depth, height, width, channels)`
+        - If `data_format="channels_first"`: 5D tensor
+            `(batch_size, channels, depth, height, width)`
+
+    Output shape:
+        - If `data_format="channels_last"`:
+            `(batch_size, output_depth, output_height, output_width, channels)`
+        - If `data_format="channels_first"`:
+            `(batch_size, channels, output_depth, output_height, output_width)`
+
+    Examples:
+        >>> import numpy as np
+        >>> input_vol = np.random.rand(1, 32, 32, 32, 3)
+        >>> layer = AdaptiveAveragePooling3D(output_size=16)
+        >>> output_vol = layer(input_vol)
+        >>> output_vol.shape
+        (1, 16, 16, 16, 3)
+    """
+
+    def __init__(self, output_size, data_format=None, **kwargs):
+        if isinstance(output_size, int):
+            output_size_tuple = (output_size, output_size, output_size)
+        elif isinstance(output_size, (tuple, list)) and len(output_size) == 3:
+            output_size_tuple = tuple(output_size)
+        else:
+            raise TypeError(
+                f"`output_size` must be an integer or "
+                f"(depth, height, width) tuple. "
+                f"Received: {output_size} of type {type(output_size)}"
+            )
+
+        super().__init__(output_size_tuple, data_format, **kwargs)

keras/src/layers/pooling/adaptive_max_pooling1d.py

@@ -0,0 +1,65 @@
+"""Adaptive Max Pooling 1D layer."""
+
+from keras.src.api_export import keras_export
+from keras.src.layers.pooling.base_adaptive_pooling import (
+    BaseAdaptiveMaxPooling,
+)
+
+
+@keras_export("keras.layers.AdaptiveMaxPooling1D")
+class AdaptiveMaxPooling1D(BaseAdaptiveMaxPooling):
+    """Adaptive max pooling operation for 1D temporal or spatial data.
+
+    This layer applies an adaptive max pooling operation, which pools the
+    input such that the output has a target length specified by `output_size`,
+    regardless of the input length. The kernel size and stride are automatically
+    computed to achieve the target output size.
+
+    Args:
+        output_size: Integer specifying the target output length.
+        data_format: string, either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, length, channels)`.
+            `"channels_first"` corresponds to inputs with shape
+            `(batch, channels, length)`.
+            Defaults to the value found in your Keras config file at
+            `~/.keras/keras.json`. If never set, `"channels_last"` is used.
+
+    Input shape:
+        - If `data_format="channels_last"`: 3D tensor
+            `(batch_size, length, channels)`
+        - If `data_format="channels_first"`: 3D tensor
+            `(batch_size, channels, length)`
+
+    Output shape:
+        - If `data_format="channels_last"`:
+            `(batch_size, output_length, channels)`
+        - If `data_format="channels_first"`:
+            `(batch_size, channels, output_length)`
+
+    Examples:
+        >>> import numpy as np
+        >>> input_seq = np.random.rand(1, 64, 3)
+        >>> layer = AdaptiveMaxPooling1D(output_size=32)
+        >>> output_seq = layer(input_seq)
+        >>> output_seq.shape
+        (1, 32, 3)
+    """
+
+    def __init__(self, output_size, data_format=None, **kwargs):
+        if isinstance(output_size, int):
+            output_size = (output_size,)
+        elif isinstance(output_size, (tuple, list)):
+            if len(output_size) != 1:
+                raise ValueError(
+                    f"For 1D input, `output_size` tuple must have length 1. "
+                    f"Received: {output_size}"
+                )
+            output_size = tuple(output_size)
+        else:
+            raise TypeError(
+                f"`output_size` must be an integer or tuple of 1 integer. "
+                f"Received: {output_size} of type {type(output_size)}"
+            )
+
+        super().__init__(output_size, data_format, **kwargs)

keras/src/layers/pooling/adaptive_max_pooling2d.py

@@ -0,0 +1,62 @@
+"""Adaptive Max Pooling 2D layer."""
+
+from keras.src.api_export import keras_export
+from keras.src.layers.pooling.base_adaptive_pooling import (
+    BaseAdaptiveMaxPooling,
+)
+
+
+@keras_export("keras.layers.AdaptiveMaxPooling2D")
+class AdaptiveMaxPooling2D(BaseAdaptiveMaxPooling):
+    """Adaptive max pooling operation for 2D spatial data.
+
+    This layer applies an adaptive max pooling operation, which pools the
+    input such that the output has a target spatial size specified by
+    `output_size`, regardless of the input spatial size. The kernel size
+    and stride are automatically computed to achieve the target output size.
+
+    Args:
+        output_size: Integer or tuple of 2 integers specifying the
+            target output size.
+            If an integer, the same value is used for both height and width.
+        data_format: string, either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, height, width, channels)`.
+            `"channels_first"` corresponds to inputs with shape
+            `(batch, channels, height, width)`.
+            Defaults to the value found in your Keras config file at
+            `~/.keras/keras.json`. If never set, `"channels_last"` is used.
+
+    Input shape:
+        - If `data_format="channels_last"`: 4D tensor
+            `(batch_size, height, width, channels)`
+        - If `data_format="channels_first"`: 4D tensor
+            `(batch_size, channels, height, width)`
+
+    Output shape:
+        - If `data_format="channels_last"`:
+            `(batch_size, output_height, output_width, channels)`
+        - If `data_format="channels_first"`:
+            `(batch_size, channels, output_height, output_width)`
+
+    Examples:
+        >>> import numpy as np
+        >>> input_img = np.random.rand(1, 64, 64, 3)
+        >>> layer = AdaptiveMaxPooling2D(output_size=32)
+        >>> output_img = layer(input_img)
+        >>> output_img.shape
+        (1, 32, 32, 3)
+    """
+
+    def __init__(self, output_size, data_format=None, **kwargs):
+        if isinstance(output_size, int):
+            output_size_tuple = (output_size, output_size)
+        elif isinstance(output_size, (tuple, list)) and len(output_size) == 2:
+            output_size_tuple = tuple(output_size)
+        else:
+            raise TypeError(
+                f"`output_size` must be an integer or (height, width) tuple. "
+                f"Received: {output_size} of type {type(output_size)}"
+            )
+
+        super().__init__(output_size_tuple, data_format, **kwargs)

keras/src/layers/pooling/adaptive_max_pooling3d.py

@@ -0,0 +1,63 @@
+"""Adaptive Max Pooling 3D layer."""
+
+from keras.src.api_export import keras_export
+from keras.src.layers.pooling.base_adaptive_pooling import (
+    BaseAdaptiveMaxPooling,
+)
+
+
+@keras_export("keras.layers.AdaptiveMaxPooling3D")
+class AdaptiveMaxPooling3D(BaseAdaptiveMaxPooling):
+    """Adaptive max pooling operation for 3D volumetric data.
+
+    This layer applies an adaptive max pooling operation, which pools the
+    input such that the output has a target spatial size specified by
+    `output_size`, regardless of the input spatial size. The kernel size
+    and stride are automatically computed to achieve the target output size.
+
+    Args:
+        output_size: Integer or tuple of 3 integers specifying the
+            target output size.
+            If an integer, the same value is used for depth, height, and width.
+        data_format: string, either `"channels_last"` or `"channels_first"`.
+            `"channels_last"` corresponds to inputs with shape
+            `(batch, depth, height, width, channels)`.
+            `"channels_first"` corresponds to inputs with shape
+            `(batch, channels, depth, height, width)`.
+            Defaults to the value found in your Keras config file at
+            `~/.keras/keras.json`. If never set, `"channels_last"` is used.
+
+    Input shape:
+        - If `data_format="channels_last"`: 5D tensor
+            `(batch_size, depth, height, width, channels)`
+        - If `data_format="channels_first"`: 5D tensor
+            `(batch_size, channels, depth, height, width)`
+
+    Output shape:
+        - If `data_format="channels_last"`:
+            `(batch_size, output_depth, output_height, output_width, channels)`
+        - If `data_format="channels_first"`:
+            `(batch_size, channels, output_depth, output_height, output_width)`
+
+    Examples:
+        >>> import numpy as np
+        >>> input_vol = np.random.rand(1, 32, 32, 32, 3)
+        >>> layer = AdaptiveMaxPooling3D(output_size=16)
+        >>> output_vol = layer(input_vol)
+        >>> output_vol.shape
+        (1, 16, 16, 16, 3)
+    """
+
+    def __init__(self, output_size, data_format=None, **kwargs):
+        if isinstance(output_size, int):
+            output_size_tuple = (output_size, output_size, output_size)
+        elif isinstance(output_size, (tuple, list)) and len(output_size) == 3:
+            output_size_tuple = tuple(output_size)
+        else:
+            raise TypeError(
+                f"`output_size` must be an integer or "
+                f"(depth, height, width) tuple. "
+                f"Received: {output_size} of type {type(output_size)}"
+            )
+
+        super().__init__(output_size_tuple, data_format, **kwargs)

keras/src/layers/pooling/base_adaptive_pooling.py

@@ -0,0 +1,63 @@
+"""Base classes for adaptive pooling layers."""
+
+from keras.src import ops
+from keras.src.backend import config
+from keras.src.layers.layer import Layer
+
+
+class BaseAdaptivePooling(Layer):
+    """Base class shared by all adaptive pooling layers."""
+
+    def __init__(self, output_size, data_format=None, **kwargs):
+        """Initialize base adaptive pooling layer.
+
+        Args:
+            output_size: Normalized spatial output size as a tuple
+                (for example, (32,), (32, 32), or (32, 32, 32)).
+            data_format: Either "channels_last" or "channels_first".
+            **kwargs: Additional layer keyword arguments.
+        """
+        super().__init__(**kwargs)
+        self.output_size = output_size
+        self.data_format = data_format or config.image_data_format()
+        if self.data_format not in {"channels_first", "channels_last"}:
+            raise ValueError(
+                f"Invalid data_format: {self.data_format}. "
+                "Expected 'channels_first' or 'channels_last'."
+            )
+
+    def compute_output_shape(self, input_shape):
+        """Return the output shape tensor after pooling."""
+        batch_size = input_shape[0]
+        if self.data_format == "channels_last":
+            channels = input_shape[-1]
+            return (batch_size, *self.output_size, channels)
+        else:
+            channels = input_shape[1]
+            return (batch_size, channels, *self.output_size)
+
+    def get_config(self):
+        config_dict = {
+            "output_size": self.output_size,
+            "data_format": self.data_format,
+        }
+        base_config = super().get_config()
+        return {**base_config, **config_dict}
+
+
+class BaseAdaptiveAveragePooling(BaseAdaptivePooling):
+    """Base class for adaptive average pooling in 1D, 2D, and 3D."""
+
+    def call(self, inputs):
+        return ops.adaptive_average_pool(
+            inputs, output_size=self.output_size, data_format=self.data_format
+        )
+
+
+class BaseAdaptiveMaxPooling(BaseAdaptivePooling):
+    """Base class for adaptive max pooling in 1D, 2D, and 3D."""
+
+    def call(self, inputs):
+        return ops.adaptive_max_pool(
+            inputs, output_size=self.output_size, data_format=self.data_format
+        )

keras/src/layers/preprocessing/discretization.py

@@ -95,9 +95,6 @@ class Discretization(DataLayer):
         dtype=None,
         name=None,
     ):
-        if dtype is None:
-            dtype = "int64" if output_mode == "int" else backend.floatx()
-
         super().__init__(name=name, dtype=dtype)
 
         if sparse and not backend.SUPPORTS_SPARSE_TENSORS:
@@ -155,6 +152,10 @@ class Discretization(DataLayer):
     def input_dtype(self):
         return backend.floatx()
 
+    @property
+    def output_dtype(self):
+        return self.compute_dtype if self.output_mode != "int" else "int32"
+
     def adapt(self, data, steps=None):
         """Computes bin boundaries from quantiles in a input dataset.
 
@@ -213,7 +214,7 @@ class Discretization(DataLayer):
         self.summary = np.array([[], []], dtype="float32")
 
     def compute_output_spec(self, inputs):
-        return backend.KerasTensor(shape=inputs.shape, dtype=self.compute_dtype)
+        return backend.KerasTensor(shape=inputs.shape, dtype=self.output_dtype)
 
     def load_own_variables(self, store):
         if len(store) == 1:
@@ -234,7 +235,7 @@ class Discretization(DataLayer):
             indices,
             output_mode=self.output_mode,
             depth=len(self.bin_boundaries) + 1,
-            dtype=self.compute_dtype,
+            dtype=self.output_dtype,
             sparse=self.sparse,
             backend_module=self.backend,
         )

keras/src/layers/preprocessing/feature_space.py

@@ -507,10 +507,14 @@ class FeatureSpace(Layer):
 
     def adapt(self, dataset):
         if not isinstance(dataset, tf.data.Dataset):
-            raise ValueError(
-                "`adapt()` can only be called on a tf.data.Dataset. "
-                f"Received instead: {dataset} (of type {type(dataset)})"
-            )
+            if isinstance(dataset, dict):
+                dataset = tf.data.Dataset.from_tensor_slices(dataset)
+            else:
+                raise ValueError(
+                    "`adapt()` can only be called on a tf.data.Dataset or a "
+                    "dict of arrays/lists. "
+                    f"Received instead: {dataset} (of type {type(dataset)})"
+                )
 
         for name in self._list_adaptable_preprocessors():
             # Call adapt() on each individual adaptable layer.

keras/src/layers/preprocessing/image_preprocessing/aug_mix.py

@@ -316,8 +316,8 @@ class AugMix(BaseImagePreprocessingLayer):
     def get_config(self):
         config = {
             "value_range": self.value_range,
-            "num_chains": self.chain_depth,
-            "chain_depth": self.num_chains,
+            "num_chains": self.num_chains,
+            "chain_depth": self.chain_depth,
             "factor": self.factor,
             "alpha": self.alpha,
             "all_ops": self.all_ops,

keras/src/layers/preprocessing/image_preprocessing/random_contrast.py

@@ -92,8 +92,8 @@ class RandomContrast(BaseImagePreprocessingLayer):
 
     def transform_images(self, images, transformation, training=True):
         if training:
-            constrast_factor = transformation["contrast_factor"]
-            outputs = self._adjust_constrast(images, constrast_factor)
+            contrast_factor = transformation["contrast_factor"]
+            outputs = self._adjust_contrast(images, contrast_factor)
             outputs = self.backend.numpy.clip(
                 outputs, self.value_range[0], self.value_range[1]
             )
@@ -117,7 +117,7 @@ class RandomContrast(BaseImagePreprocessingLayer):
     ):
         return segmentation_masks
 
-    def _adjust_constrast(self, inputs, contrast_factor):
+    def _adjust_contrast(self, inputs, contrast_factor):
         if self.data_format == "channels_first":
             height_axis = -2
             width_axis = -1

keras/src/layers/preprocessing/image_preprocessing/resizing.py

@@ -66,6 +66,16 @@ class Resizing(BaseImagePreprocessingLayer):
             `~/.keras/keras.json`. If you never set it, then it will be
             `"channels_last"`.
         **kwargs: Base layer keyword arguments, such as `name` and `dtype`.
+
+    Example:
+
+    ```python
+    (x_train, y_train), _ = keras.datasets.cifar10.load_data()
+    image = x_train[0]
+    resizer = keras.layers.Resizing(128, 128)
+    resized_image = resizer(image)
+    print("original:", image.shape, "resized:", resized_image.shape)
+    ```
     """
 
     _USE_BASE_FACTOR = False

keras/src/layers/preprocessing/index_lookup.py

@@ -4,6 +4,7 @@ import numpy as np
 
 from keras.src import backend
 from keras.src.layers.layer import Layer
+from keras.src.saving import serialization_lib
 from keras.src.utils import argument_validation
 from keras.src.utils import numerical_utils
 from keras.src.utils import tf_utils
@@ -178,7 +179,12 @@ class IndexLookup(Layer):
         self.vocabulary_dtype = tf.as_dtype(vocabulary_dtype).name
         self._frozen_vocab_size = kwargs.pop("vocabulary_size", None)
 
-        self.input_vocabulary = vocabulary
+        # Remember original `vocabulary` as `input_vocabulary` for serialization
+        # via `get_config`. However, if `vocabulary` is a file path or a URL, we
+        # serialize the vocabulary as an asset and clear the original path/URL.
+        self.input_vocabulary = (
+            vocabulary if not isinstance(vocabulary, str) else None
+        )
         self.input_idf_weights = idf_weights
 
         # We set this hidden attr to
@@ -382,6 +388,18 @@ class IndexLookup(Layer):
             )
 
         if isinstance(vocabulary, str):
+            if serialization_lib.in_safe_mode():
+                raise ValueError(
+                    "Requested the loading of a vocabulary file outside of the "
+                    "model archive. This carries a potential risk of loading "
+                    "arbitrary and sensitive files and thus it is disallowed "
+                    "by default. If you trust the source of the artifact, you "
+                    "can override this error by passing `safe_mode=False` to "
+                    "the loading function, or calling "
+                    "`keras.config.enable_unsafe_deserialization(). "
+                    f"Vocabulary file: '{vocabulary}'"
+                )
+
             if not tf.io.gfile.exists(vocabulary):
                 raise ValueError(
                     f"Vocabulary file {vocabulary} does not exist."

keras/src/layers/preprocessing/normalization.py

@@ -6,6 +6,7 @@ from keras.src import backend
 from keras.src import ops
 from keras.src.api_export import keras_export
 from keras.src.layers.preprocessing.data_layer import DataLayer
+from keras.src.trainers.data_adapters.py_dataset_adapter import PyDataset
 from keras.src.utils.module_utils import tensorflow as tf
 
 
@@ -229,6 +230,18 @@ class Normalization(DataLayer):
                 # Batch dataset if it isn't batched
                 data = data.batch(128)
             input_shape = tuple(data.element_spec.shape)
+        elif isinstance(data, PyDataset):
+            data = data[0]
+            if isinstance(data, tuple):
+                # handling (x, y) or (x, y, sample_weight)
+                data = data[0]
+            input_shape = data.shape
+        else:
+            raise TypeError(
+                f"Unsupported data type: {type(data)}. `adapt` supports "
+                f"`np.ndarray`, backend tensors, `tf.data.Dataset`, and "
+                f"`keras.utils.PyDataset`."
+            )
 
         if not self.built:
             self.build(input_shape)
@@ -248,7 +261,7 @@ class Normalization(DataLayer):
         elif backend.is_tensor(data):
             total_mean = ops.mean(data, axis=self._reduce_axis)
             total_var = ops.var(data, axis=self._reduce_axis)
-        elif isinstance(data, tf.data.Dataset):
+        elif isinstance(data, (tf.data.Dataset, PyDataset)):
             total_mean = ops.zeros(self._mean_and_var_shape)
             total_var = ops.zeros(self._mean_and_var_shape)
             total_count = 0