keras-nightly 3.12.0.dev2025082003__py3-none-any.whl → 3.12.0.dev2025082203__py3-none-any.whl

keras/src/datasets/reuters.py

@@ -87,7 +87,7 @@ def load_data(
     )
     path = get_file(
         fname=path,
-        origin=origin_folder + "reuters.npz",
+        origin=f"{origin_folder}reuters.npz",
         file_hash=(  # noqa: E501
             "d6586e694ee56d7a4e65172e12b3e987c03096cb01eab99753921ef915959916"
         ),
@@ -156,7 +156,7 @@ def get_word_index(path="reuters_word_index.json"):
     )
     path = get_file(
         path,
-        origin=origin_folder + "reuters_word_index.json",
+        origin=f"{origin_folder}reuters_word_index.json",
         file_hash="4d44cc38712099c9e383dc6e5f11a921",
     )
     with open(path) as f:

keras/src/dtype_policies/dtype_policy.py

@@ -3,7 +3,7 @@ from keras.src import ops
 from keras.src.api_export import keras_export
 from keras.src.backend.common import global_state
 
-QUANTIZATION_MODES = ("int8", "float8", "int4")
+QUANTIZATION_MODES = ("int8", "float8", "int4", "gptq")
 
 
 @keras_export(

keras/src/dtype_policies/dtype_policy_map.py

@@ -74,7 +74,7 @@ class DTypePolicyMap(DTypePolicy, MutableMapping):
 
     @property
     def name(self):
-        return "map_" + self.default_policy._name
+        return f"map_{self.default_policy._name}"
 
     @property
     def default_policy(self):

keras/src/export/tf2onnx_lib.py

@@ -157,9 +157,7 @@ def patch_tf2onnx():
         ):
             a = copy.deepcopy(a)
             tensor_name = (
-                self.name.strip()
-                + "_"
-                + str(external_tensor_storage.name_counter)
+                f"{self.name.strip()}_{external_tensor_storage.name_counter}"
             )
             for c in '~"#%&*:<>?/\\{|}':
                 tensor_name = tensor_name.replace(c, "_")

keras/src/layers/attention/attention.py

@@ -176,6 +176,8 @@ class Attention(Layer):
             # Bias so padding positions do not contribute to attention
             # distribution.  Note 65504. is the max float16 value.
             max_value = 65504.0 if scores.dtype == "float16" else 1.0e9
+            if len(padding_mask.shape) == 2:
+                padding_mask = ops.expand_dims(padding_mask, axis=-2)
             scores -= max_value * ops.cast(padding_mask, dtype=scores.dtype)
 
         weights = ops.softmax(scores, axis=-1)

keras/src/layers/core/lambda_layer.py

@@ -167,14 +167,15 @@ class Lambda(Layer):
         )
 
     @staticmethod
-    def _raise_for_lambda_deserialization(arg_name, safe_mode):
+    def _raise_for_lambda_deserialization(safe_mode):
         if safe_mode:
             raise ValueError(
-                f"The `{arg_name}` of this `Lambda` layer is a Python lambda. "
-                "Deserializing it is unsafe. If you trust the source of the "
-                "config artifact, you can override this error "
-                "by passing `safe_mode=False` "
-                "to `from_config()`, or calling "
+                "Requested the deserialization of a `Lambda` layer whose "
+                "`function` is a Python lambda. This carries a potential risk "
+                "of arbitrary code execution 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()."
             )
 
@@ -187,7 +188,7 @@ class Lambda(Layer):
             and "class_name" in fn_config
             and fn_config["class_name"] == "__lambda__"
         ):
-            cls._raise_for_lambda_deserialization("function", safe_mode)
+            cls._raise_for_lambda_deserialization(safe_mode)
             inner_config = fn_config["config"]
             fn = python_utils.func_load(
                 inner_config["code"],
@@ -206,7 +207,7 @@ class Lambda(Layer):
                 and "class_name" in fn_config
                 and fn_config["class_name"] == "__lambda__"
             ):
-                cls._raise_for_lambda_deserialization("function", safe_mode)
+                cls._raise_for_lambda_deserialization(safe_mode)
                 inner_config = fn_config["config"]
                 fn = python_utils.func_load(
                     inner_config["code"],

keras/src/layers/input_spec.py

@@ -94,12 +94,12 @@ class InputSpec:
 
     def __repr__(self):
         spec = [
-            ("dtype=" + str(self.dtype)) if self.dtype else "",
-            ("shape=" + str(self.shape)) if self.shape else "",
-            ("ndim=" + str(self.ndim)) if self.ndim else "",
-            ("max_ndim=" + str(self.max_ndim)) if self.max_ndim else "",
-            ("min_ndim=" + str(self.min_ndim)) if self.min_ndim else "",
-            ("axes=" + str(self.axes)) if self.axes else "",
+            (f"dtype={str(self.dtype)}") if self.dtype else "",
+            (f"shape={str(self.shape)}") if self.shape else "",
+            (f"ndim={str(self.ndim)}") if self.ndim else "",
+            (f"max_ndim={str(self.max_ndim)}") if self.max_ndim else "",
+            (f"min_ndim={str(self.min_ndim)}") if self.min_ndim else "",
+            (f"axes={str(self.axes)}") if self.axes else "",
         ]
         return f"InputSpec({', '.join(x for x in spec if x)})"
 

keras/src/layers/layer.py

@@ -1337,7 +1337,7 @@ class Layer(BackendLayer, Operation):
         else:
             attr_name = str(attr)
             attr_type = "attribute"
-        msg = " " + msg if msg is not None else ""
+        msg = f" {msg}" if msg is not None else ""
         return NotImplementedError(
             f"Layer {self.__class__.__name__} does not have a `{attr_name}` "
             f"{attr_type} implemented.{msg}"

keras/src/layers/preprocessing/category_encoding.py

@@ -1,12 +1,12 @@
 from keras.src.api_export import keras_export
 from keras.src.backend import KerasTensor
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 from keras.src.utils import backend_utils
 from keras.src.utils import numerical_utils
 
 
 @keras_export("keras.layers.CategoryEncoding")
-class CategoryEncoding(TFDataLayer):
+class CategoryEncoding(DataLayer):
     """A preprocessing layer which encodes integer features.
 
     This layer provides options for condensing data into a categorical encoding
@@ -15,7 +15,7 @@ class CategoryEncoding(TFDataLayer):
     inputs. For integer inputs where the total number of tokens is not known,
     use `keras.layers.IntegerLookup` instead.
 
-    **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).
 
     Examples:

keras/src/layers/preprocessing/data_layer.py

@@ -0,0 +1,159 @@
+import keras.src.backend
+from keras.src import tree
+from keras.src.layers.layer import Layer
+from keras.src.random.seed_generator import SeedGenerator
+from keras.src.utils import backend_utils
+from keras.src.utils import jax_utils
+from keras.src.utils import tracking
+
+
+class DataLayer(Layer):
+    """Layer designed for safe use in `tf.data` or `grain` pipeline.
+
+    This layer overrides the `__call__` method to ensure that the correct
+    backend is used and that computation is performed on the CPU.
+
+    The `call()` method in subclasses should use `self.backend` ops. If
+    randomness is needed, define both `seed` and `generator` in `__init__` and
+    retrieve the running seed using `self._get_seed_generator()`. If the layer
+    has weights in `__init__` or `build()`, use `convert_weight()` to ensure
+    they are in the correct backend.
+
+    **Note:** This layer and its subclasses only support a single input tensor.
+
+    Examples:
+
+    **Custom `DataLayer` subclass:**
+
+    ```python
+    from keras.src.layers.preprocessing.data_layer import DataLayer
+    from keras.src.random import SeedGenerator
+
+
+    class BiasedRandomRGBToHSVLayer(DataLayer):
+        def __init__(self, seed=None, **kwargs):
+            super().__init__(**kwargs)
+            self.probability_bias = ops.convert_to_tensor(0.01)
+            self.seed = seed
+            self.generator = SeedGenerator(seed)
+
+        def call(self, inputs):
+            images_shape = self.backend.shape(inputs)
+            batch_size = 1 if len(images_shape) == 3 else images_shape[0]
+            seed = self._get_seed_generator(self.backend._backend)
+
+            probability = self.backend.random.uniform(
+                shape=(batch_size,),
+                minval=0.0,
+                maxval=1.0,
+                seed=seed,
+            )
+            probability = self.backend.numpy.add(
+                probability, self.convert_weight(self.probability_bias)
+            )
+            hsv_images = self.backend.image.rgb_to_hsv(inputs)
+            return self.backend.numpy.where(
+                probability[:, None, None, None] > 0.5,
+                hsv_images,
+                inputs,
+            )
+
+        def compute_output_shape(self, input_shape):
+            return input_shape
+    ```
+
+    **Using as a regular Keras layer:**
+
+    ```python
+    import numpy as np
+
+    x = np.random.uniform(size=(1, 16, 16, 3)).astype("float32")
+    print(BiasedRandomRGBToHSVLayer()(x).shape)  # (1, 16, 16, 3)
+    ```
+
+    **Using in a `tf.data` pipeline:**
+
+    ```python
+    import tensorflow as tf
+
+    tf_ds = tf.data.Dataset.from_tensors(x)
+    tf_ds = tf_ds.map(BiasedRandomRGBToHSVLayer())
+    print([x.shape for x in tf_ds])  # [(1, 16, 16, 3)]
+    ```
+
+    **Using in a `grain` pipeline:**
+
+    ```python
+    import grain
+
+    grain_ds = grain.MapDataset.source([x])
+    grain_ds = grain_ds.map(BiasedRandomRGBToHSVLayer())
+    print([x.shape for x in grain_ds])  # [(1, 16, 16, 3)]
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.backend = backend_utils.DynamicBackend()
+        self._allow_non_tensor_positional_args = True
+
+    def __call__(self, inputs, **kwargs):
+        sample_input = tree.flatten(inputs)[0]
+        if (
+            not isinstance(sample_input, keras.KerasTensor)
+            and backend_utils.in_tf_graph()
+            and not jax_utils.is_in_jax_tracing_scope(sample_input)
+        ):
+            # We're in a TF graph, e.g. a tf.data pipeline.
+            self.backend.set_backend("tensorflow")
+            inputs = tree.map_structure(
+                lambda x: self.backend.convert_to_tensor(
+                    x, dtype=self.compute_dtype
+                ),
+                inputs,
+            )
+            switch_convert_input_args = False
+            if self._convert_input_args:
+                self._convert_input_args = False
+                switch_convert_input_args = True
+            try:
+                outputs = super().__call__(inputs, **kwargs)
+            finally:
+                self.backend.reset()
+                if switch_convert_input_args:
+                    self._convert_input_args = True
+            return outputs
+        elif (
+            not isinstance(sample_input, keras.KerasTensor)
+            and backend_utils.in_grain_data_pipeline()
+        ):
+            # We're in a Grain data pipeline. Force computation and data
+            # placement to CPU.
+            with keras.src.backend.device_scope("cpu"):
+                return super().__call__(inputs, **kwargs)
+        else:
+            return super().__call__(inputs, **kwargs)
+
+    @tracking.no_automatic_dependency_tracking
+    def _get_seed_generator(self, backend=None):
+        if not hasattr(self, "seed") or not hasattr(self, "generator"):
+            raise ValueError(
+                "The `seed` and `generator` variable must be set in the "
+                "`__init__` method before calling `_get_seed_generator()`."
+            )
+        if backend is None or backend == keras.backend.backend():
+            return self.generator
+        if not hasattr(self, "_backend_generators"):
+            self._backend_generators = {}
+        if backend in self._backend_generators:
+            return self._backend_generators[backend]
+        seed_generator = SeedGenerator(self.seed, backend=self.backend)
+        self._backend_generators[backend] = seed_generator
+        return seed_generator
+
+    def convert_weight(self, weight):
+        """Convert the weight if it is from the a different backend."""
+        if self.backend.name == keras.backend.backend():
+            return weight
+        else:
+            weight = keras.ops.convert_to_numpy(weight)
+            return self.backend.convert_to_tensor(weight)

keras/src/layers/preprocessing/discretization.py

@@ -2,21 +2,21 @@ import numpy as np
 
 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.utils import argument_validation
 from keras.src.utils import numerical_utils
 from keras.src.utils.module_utils import tensorflow as tf
 
 
 @keras_export("keras.layers.Discretization")
-class Discretization(TFDataLayer):
+class Discretization(DataLayer):
     """A preprocessing layer which buckets continuous features by ranges.
 
     This layer will place each element of its input data into one of several
     contiguous ranges and output an integer index indicating which range each
     element was placed in.
 
-    **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/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: