keras-nightly 3.12.0.dev2025082103__py3-none-any.whl → 3.12.0.dev2025082303__py3-none-any.whl

keras/src/layers/preprocessing/feature_space.py

@@ -3,7 +3,7 @@ from keras.src import layers
 from keras.src import tree
 from keras.src.api_export import keras_export
 from keras.src.layers.layer import Layer
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 from keras.src.saving import saving_lib
 from keras.src.saving import serialization_lib
 from keras.src.saving.keras_saveable import KerasSaveable
@@ -723,7 +723,7 @@ class FeatureSpace(Layer):
                 data[name] = tf.expand_dims(x, -1)
 
         with backend_utils.TFGraphScope():
-            # This scope is to make sure that inner TFDataLayers
+            # This scope is to make sure that inner DataLayers
             # will not convert outputs back to backend-native --
             # they should be TF tensors throughout
             preprocessed_data = self._preprocess_features(data)
@@ -808,7 +808,7 @@ class FeatureSpace(Layer):
         return
 
 
-class TFDConcat(TFDataLayer):
+class TFDConcat(DataLayer):
     def __init__(self, axis, **kwargs):
         super().__init__(**kwargs)
         self.axis = axis
@@ -817,6 +817,6 @@ class TFDConcat(TFDataLayer):
         return self.backend.numpy.concatenate(xs, axis=self.axis)
 
 
-class TFDIdentity(TFDataLayer):
+class TFDIdentity(DataLayer):
     def call(self, x):
         return x

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

@@ -43,6 +43,13 @@ class AugMix(BaseImagePreprocessingLayer):
     in num_chains different ways, with each chain consisting of
     chain_depth augmentations.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
+    References:
+        - [AugMix paper](https://arxiv.org/pdf/1912.02781)
+        - [Official Code](https://github.com/google-research/augmix)
+
     Args:
         value_range: the range of values the incoming images will have.
             Represented as a two number tuple written (low, high).
@@ -64,10 +71,6 @@ class AugMix(BaseImagePreprocessingLayer):
         interpolation: The interpolation method to use for resizing operations.
             Options include `"nearest"`, `"bilinear"`. Default is `"bilinear"`.
         seed: Integer. Used to create a random seed.
-
-    References:
-        - [AugMix paper](https://arxiv.org/pdf/1912.02781)
-        - [Official Code](https://github.com/google-research/augmix)
     """
 
     _USE_BASE_FACTOR = False

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

@@ -17,6 +17,9 @@ class AutoContrast(BaseImagePreprocessingLayer):
 
     This layer is active at both training and inference time.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         value_range: Range of values the incoming images will have.
             Represented as a two number tuple written `(low, high)`.

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

@@ -1,13 +1,13 @@
 import math
 
 from keras.src.backend import config as backend_config
+from keras.src.layers.preprocessing.data_layer import DataLayer
 from keras.src.layers.preprocessing.image_preprocessing.bounding_boxes.validation import (  # noqa: E501
     densify_bounding_boxes,
 )
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
 
 
-class BaseImagePreprocessingLayer(TFDataLayer):
+class BaseImagePreprocessingLayer(DataLayer):
     _USE_BASE_FACTOR = True
     _FACTOR_BOUNDS = (-1, 1)
 

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

@@ -36,7 +36,7 @@ class CenterCrop(BaseImagePreprocessingLayer):
     If the input height/width is even and the target height/width is odd (or
     inversely), the input image is left-padded by 1 pixel.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Args:

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

@@ -13,6 +13,12 @@ class CutMix(BaseImagePreprocessingLayer):
     between two images in the dataset, while the labels are also mixed
     proportionally to the area of the patches.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
+    References:
+       - [CutMix paper]( https://arxiv.org/abs/1905.04899).
+
     Args:
         factor: A single float or a tuple of two floats between 0 and 1.
             If a tuple of numbers is passed, a `factor` is sampled
@@ -23,9 +29,6 @@ class CutMix(BaseImagePreprocessingLayer):
             in patch sizes, leading to more diverse and larger mixed patches.
             Defaults to 1.
         seed: Integer. Used to create a random seed.
-
-    References:
-       - [CutMix paper]( https://arxiv.org/abs/1905.04899).
     """
 
     _USE_BASE_FACTOR = False

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

@@ -18,7 +18,7 @@ class Equalization(BaseImagePreprocessingLayer):
     equalization independently on each color channel. At inference time,
     the equalization is consistently applied.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Args:

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

@@ -8,6 +8,9 @@ from keras.src.layers.preprocessing.image_preprocessing.base_image_preprocessing
 class MaxNumBoundingBoxes(BaseImagePreprocessingLayer):
     """Ensure the maximum number of bounding boxes.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         max_number: Desired output number of bounding boxes.
         padding_value: The padding value of the `boxes` and `labels` in

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

@@ -11,6 +11,13 @@ from keras.src.utils import backend_utils
 class MixUp(BaseImagePreprocessingLayer):
     """MixUp implements the MixUp data augmentation technique.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
+    References:
+        - [MixUp paper](https://arxiv.org/abs/1710.09412).
+        - [MixUp for Object Detection paper](https://arxiv.org/pdf/1902.04103).
+
     Args:
         alpha: Float between 0 and 1. Controls the blending strength.
                Smaller values mean less mixing, while larger values allow
@@ -18,10 +25,6 @@ class MixUp(BaseImagePreprocessingLayer):
                recommended for ImageNet1k classification.
         seed: Integer. Used to create a random seed.
 
-    References:
-        - [MixUp paper](https://arxiv.org/abs/1710.09412).
-        - [MixUp for Object Detection paper](https://arxiv.org/pdf/1902.04103).
-
     Example:
     ```python
     (images, labels), _ = keras.datasets.cifar10.load_data()

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

@@ -15,6 +15,9 @@ class RandAugment(BaseImagePreprocessingLayer):
     policy implemented by this layer has been benchmarked extensively and is
     effective on a wide variety of datasets.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     References:
         - [RandAugment](https://arxiv.org/abs/1909.13719)
 
@@ -29,7 +32,6 @@ class RandAugment(BaseImagePreprocessingLayer):
         interpolation: The interpolation method to use for resizing operations.
             Options include `nearest`, `bilinear`. Default is `bilinear`.
         seed: Integer. Used to create a random seed.
-
     """
 
     _USE_BASE_FACTOR = False

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

@@ -13,7 +13,7 @@ class RandomBrightness(BaseImagePreprocessingLayer):
     images. At inference time, the output will be identical to the input.
     Call the layer with `training=True` to adjust the brightness of the input.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Args:

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

@@ -13,6 +13,9 @@ class RandomColorDegeneration(BaseImagePreprocessingLayer):
     color. It then takes a weighted average between original image and the
     degenerated image. This makes colors appear more dull.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         factor: A tuple of two floats or a single float.
             `factor` controls the extent to which the

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

@@ -16,6 +16,9 @@ class RandomColorJitter(BaseImagePreprocessingLayer):
     and hue image processing operation sequentially and randomly on the
     input.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         value_range: the range of values the incoming images will have.
             Represented as a two number tuple written [low, high].

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

@@ -21,7 +21,7 @@ class RandomContrast(BaseImagePreprocessingLayer):
     in integer or floating point dtype.
     By default, the layer will output floats.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Input shape:

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

@@ -30,7 +30,7 @@ class RandomCrop(BaseImagePreprocessingLayer):
     of integer or floating point dtype. By default, the layer will output
     floats.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Input shape:

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

@@ -14,6 +14,9 @@ class RandomElasticTransform(BaseImagePreprocessingLayer):
     distortion is controlled by the `scale` parameter, while the `factor`
     determines the probability of applying the transformation.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         factor: A single float or a tuple of two floats.
             `factor` controls the probability of applying the transformation.

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

@@ -13,6 +13,12 @@ class RandomErasing(BaseImagePreprocessingLayer):
     an image are erased (replaced by a constant value or noise)
     during training to improve generalization.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
+    References:
+       - [Random Erasing paper](https://arxiv.org/abs/1708.04896).
+
     Args:
         factor: A single float or a tuple of two floats.
             `factor` controls the probability of applying the transformation.
@@ -35,9 +41,6 @@ class RandomErasing(BaseImagePreprocessingLayer):
             typically either `[0, 1]` or `[0, 255]` depending on how your
             preprocessing pipeline is set up.
         seed: Integer. Used to create a random seed.
-
-    References:
-       - [Random Erasing paper](https://arxiv.org/abs/1708.04896).
     """
 
     _USE_BASE_FACTOR = False

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

@@ -27,7 +27,7 @@ class RandomFlip(BaseImagePreprocessingLayer):
     of integer or floating point dtype.
     By default, the layer will output floats.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Input shape:

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

@@ -13,6 +13,9 @@ class RandomGaussianBlur(BaseImagePreprocessingLayer):
     randomly selected degree of blurring, controlled by the `factor` and
     `sigma` arguments.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         factor: A single float or a tuple of two floats.
             `factor` controls the extent to which the image hue is impacted.

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

@@ -19,7 +19,7 @@ class RandomGrayscale(BaseImagePreprocessingLayer):
     image using standard RGB to grayscale conversion coefficients. Images
     that are not selected for conversion remain unchanged.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Args:

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

@@ -14,6 +14,9 @@ class RandomHue(BaseImagePreprocessingLayer):
     The image hue is adjusted by converting the image(s) to HSV and rotating the
     hue channel (H) by delta. The image is then converted back to RGB.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         factor: A single float or a tuple of two floats.
             `factor` controls the extent to which the

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

@@ -14,6 +14,9 @@ class RandomInvert(BaseImagePreprocessingLayer):
     complementary values. Images that are not selected for inversion
     remain unchanged.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         factor: A single float or a tuple of two floats.
             `factor` controls the probability of inverting the image colors.

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

@@ -20,6 +20,9 @@ class RandomPerspective(BaseImagePreprocessingLayer):
     corner points, simulating a 3D-like transformation. The amount of distortion
     is controlled by the `factor` and `scale` parameters.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         factor: A float or a tuple of two floats.
             Represents the probability of applying the perspective

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

@@ -8,6 +8,9 @@ from keras.src.layers.preprocessing.image_preprocessing.base_image_preprocessing
 class RandomPosterization(BaseImagePreprocessingLayer):
     """Reduces the number of bits for each color channel.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     References:
     - [AutoAugment: Learning Augmentation Policies from Data](https://arxiv.org/abs/1805.09501)
     - [RandAugment: Practical automated data augmentation with a reduced search space](https://arxiv.org/abs/1909.13719)

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

@@ -23,7 +23,7 @@ class RandomRotation(BaseImagePreprocessingLayer):
     of integer or floating point dtype.
     By default, the layer will output floats.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Input shape:

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

@@ -13,6 +13,9 @@ class RandomSaturation(BaseImagePreprocessingLayer):
     This layer will randomly increase/reduce the saturation for the input RGB
     images.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         factor: A tuple of two floats or a single float.
             `factor` controls the extent to which the image saturation

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

@@ -13,6 +13,9 @@ class RandomSharpness(BaseImagePreprocessingLayer):
     original image and the processed image. This operation adjusts the clarity
     of the edges in an image, ranging from blurred to enhanced sharpness.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         factor: A tuple of two floats or a single float.
             `factor` controls the extent to which the image sharpness

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

@@ -23,6 +23,9 @@ class RandomShear(BaseImagePreprocessingLayer):
     regions created during the transformation are filled according to the
     `fill_mode` and `fill_value` parameters.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         x_factor: A tuple of two floats. For each augmented image, a value
             is sampled from the provided range. If a float is passed, the

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

@@ -23,6 +23,9 @@ class RandomTranslation(BaseImagePreprocessingLayer):
     of integer or floating point dtype. By default, the layer will output
     floats.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Input shape:
         3D (unbatched) or 4D (batched) tensor with shape:
         `(..., height, width, channels)`, in `"channels_last"` format,
@@ -34,9 +37,6 @@ class RandomTranslation(BaseImagePreprocessingLayer):
         or `(..., channels, target_height, target_width)`,
         in `"channels_first"` format.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
-    (independently of which backend you're using).
-
     Args:
         height_factor: a float represented as fraction of value, or a tuple of
             size 2 representing lower and upper bound for shifting vertically. A

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

@@ -24,6 +24,9 @@ class RandomZoom(BaseImagePreprocessingLayer):
     of integer or floating point dtype.
     By default, the layer will output floats.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Input shape:
         3D (unbatched) or 4D (batched) tensor with shape:
         `(..., height, width, channels)`, in `"channels_last"` format,
@@ -35,9 +38,6 @@ class RandomZoom(BaseImagePreprocessingLayer):
         or `(..., channels, target_height, target_width)`,
         in `"channels_first"` format.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
-    (independently of which backend you're using).
-
     Args:
         height_factor: a float represented as fraction of value, or a tuple of
             size 2 representing lower and upper bound for zooming vertically.

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

@@ -21,6 +21,9 @@ class Resizing(BaseImagePreprocessingLayer):
     format. Input pixel values can be of any range
     (e.g. `[0., 1.)` or `[0, 255]`).
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Input shape:
         3D (unbatched) or 4D (batched) tensor with shape:
         `(..., height, width, channels)`, in `"channels_last"` format,
@@ -32,9 +35,6 @@ class Resizing(BaseImagePreprocessingLayer):
         or `(..., channels, target_height, target_width)`,
         in `"channels_first"` format.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
-    (independently of which backend you're using).
-
     Args:
         height: Integer, the height of the output shape.
         width: Integer, the width of the output shape.

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

@@ -15,6 +15,9 @@ class Solarization(BaseImagePreprocessingLayer):
     to all values. When created with specified `threshold` the layer only
     augments pixels that are above the `threshold` value.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         addition_factor: (Optional)  A tuple of two floats or a single float,
             between 0 and 1.

keras/src/layers/preprocessing/mel_spectrogram.py

@@ -1,5 +1,5 @@
 from keras.src.api_export import keras_export
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 
 # mel spectrum constants.
 _MEL_BREAK_FREQUENCY_HERTZ = 700.0
@@ -7,7 +7,7 @@ _MEL_HIGH_FREQUENCY_Q = 1127.0
 
 
 @keras_export("keras.layers.MelSpectrogram")
-class MelSpectrogram(TFDataLayer):
+class MelSpectrogram(DataLayer):
     """A preprocessing layer to convert raw audio signals to Mel spectrograms.
 
     This layer takes `float32`/`float64` single or batched audio signal as
@@ -24,10 +24,37 @@ class MelSpectrogram(TFDataLayer):
     speech and music processing tasks like speech recognition, speaker
     identification, and music genre classification.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     References:
     - [Spectrogram](https://en.wikipedia.org/wiki/Spectrogram),
     - [Mel scale](https://en.wikipedia.org/wiki/Mel_scale).
 
+    Args:
+        fft_length: Integer, size of the FFT window.
+        sequence_stride: Integer, number of samples between successive STFT
+            columns.
+        sequence_length: Integer, size of the window used for applying
+            `window` to each audio frame. If `None`, defaults to `fft_length`.
+        window: String, name of the window function to use. Available values
+            are `"hann"` and `"hamming"`. If `window` is a tensor, it will be
+            used directly as the window and its length must be
+            `sequence_length`. If `window` is `None`, no windowing is
+            used. Defaults to `"hann"`.
+        sampling_rate: Integer, sample rate of the input signal.
+        num_mel_bins: Integer, number of mel bins to generate.
+        min_freq: Float, minimum frequency of the mel bins.
+        max_freq: Float, maximum frequency of the mel bins.
+            If `None`, defaults to `sampling_rate / 2`.
+        power_to_db: If True, convert the power spectrogram to decibels.
+        top_db: Float, minimum negative cut-off `max(10 * log10(S)) - top_db`.
+        mag_exp: Float, exponent for the magnitude spectrogram.
+            1 for magnitude, 2 for power, etc. Default is 2.
+        ref_power: Float, the power is scaled relative to it
+            `10 * log10(S / ref_power)`.
+        min_power: Float, minimum value for power and `ref_power`.
+
     Examples:
 
     **Unbatched audio signal**
@@ -55,29 +82,6 @@ class MelSpectrogram(TFDataLayer):
         2D (unbatched) or 3D (batched) tensor with
         shape:`(..., num_mel_bins, time)`.
 
-    Args:
-        fft_length: Integer, size of the FFT window.
-        sequence_stride: Integer, number of samples between successive STFT
-            columns.
-        sequence_length: Integer, size of the window used for applying
-            `window` to each audio frame. If `None`, defaults to `fft_length`.
-        window: String, name of the window function to use. Available values
-            are `"hann"` and `"hamming"`. If `window` is a tensor, it will be
-            used directly as the window and its length must be
-            `sequence_length`. If `window` is `None`, no windowing is
-            used. Defaults to `"hann"`.
-        sampling_rate: Integer, sample rate of the input signal.
-        num_mel_bins: Integer, number of mel bins to generate.
-        min_freq: Float, minimum frequency of the mel bins.
-        max_freq: Float, maximum frequency of the mel bins.
-            If `None`, defaults to `sampling_rate / 2`.
-        power_to_db: If True, convert the power spectrogram to decibels.
-        top_db: Float, minimum negative cut-off `max(10 * log10(S)) - top_db`.
-        mag_exp: Float, exponent for the magnitude spectrogram.
-            1 for magnitude, 2 for power, etc. Default is 2.
-        ref_power: Float, the power is scaled relative to it
-            `10 * log10(S / ref_power)`.
-        min_power: Float, minimum value for power and `ref_power`.
     """
 
     def __init__(

keras/src/layers/preprocessing/normalization.py

@@ -5,12 +5,12 @@ import numpy as np
 from keras.src import backend
 from keras.src import ops
 from keras.src.api_export import keras_export
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 from keras.src.utils.module_utils import tensorflow as tf
 
 
 @keras_export("keras.layers.Normalization")
-class Normalization(TFDataLayer):
+class Normalization(DataLayer):
     """A preprocessing layer that normalizes continuous features.
 
     This layer will shift and scale inputs into a distribution centered around
@@ -23,6 +23,9 @@ class Normalization(TFDataLayer):
     variance of the data and store them as the layer's weights. `adapt()` should
     be called before `fit()`, `evaluate()`, or `predict()`.
 
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
+    (independently of which backend you're using).
+
     Args:
         axis: Integer, tuple of integers, or None. The axis or axes that should
             have a separate mean and variance for each index in the shape.

keras/src/layers/preprocessing/rescaling.py

@@ -1,11 +1,11 @@
 from keras.src import backend
 from keras.src.api_export import keras_export
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 from keras.src.saving import serialization_lib
 
 
 @keras_export("keras.layers.Rescaling")
-class Rescaling(TFDataLayer):
+class Rescaling(DataLayer):
     """A preprocessing layer which rescales input values to a new range.
 
     This layer rescales every value of an input (often an image) by multiplying
@@ -23,7 +23,7 @@ class Rescaling(TFDataLayer):
     of integer or floating point dtype, and by default the layer will output
     floats.
 
-    **Note:** This layer is safe to use inside a `tf.data` pipeline
+    **Note:** This layer is safe to use inside a `tf.data` or `grain` pipeline
     (independently of which backend you're using).
 
     Args:

keras/src/layers/rnn/bidirectional.py

@@ -109,16 +109,16 @@ class Bidirectional(Layer):
         # Recreate the forward layer from the original layer config, so that it
         # will not carry over any state from the layer.
         config = serialization_lib.serialize_keras_object(layer)
-        config["config"]["name"] = "forward_" + utils.removeprefix(
-            layer.name, "forward_"
+        config["config"]["name"] = (
+            f"forward_{utils.removeprefix(layer.name, 'forward_')}"
         )
         self.forward_layer = serialization_lib.deserialize_keras_object(config)
 
         if backward_layer is None:
             config = serialization_lib.serialize_keras_object(layer)
             config["config"]["go_backwards"] = True
-            config["config"]["name"] = "backward_" + utils.removeprefix(
-                layer.name, "backward_"
+            config["config"]["name"] = (
+                f"backward_{utils.removeprefix(layer.name, 'backward_')}"
             )
             self.backward_layer = serialization_lib.deserialize_keras_object(
                 config