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

keras/src/backend/tensorflow/trainer.py

@@ -1,4 +1,5 @@
 import contextlib
+import functools
 import warnings
 
 import numpy as np
@@ -107,6 +108,21 @@ class TensorFlowTrainer(base_trainer.Trainer):
             y_pred = self(x)
         return y_pred
 
+    def _autoconvert_optionals(self, step_func):
+        # Wrapper converting (nested) TF Optional in input data to None
+        @functools.wraps(step_func)
+        def wrapper(data):
+            converted_data = tree.map_structure(
+                lambda i: (
+                    None if isinstance(i, tf.experimental.Optional) else i
+                ),
+                data,
+            )
+            result = step_func(converted_data)
+            return result
+
+        return wrapper
+
     def _make_function(self, step_function):
         @tf.autograph.experimental.do_not_convert
         def one_step_on_data(data):
@@ -125,6 +141,7 @@ class TensorFlowTrainer(base_trainer.Trainer):
                 reduce_retracing=True,
                 jit_compile=self.jit_compile,
             )
+        one_step_on_data = self._autoconvert_optionals(one_step_on_data)
 
         @tf.autograph.experimental.do_not_convert
         def multi_step_on_iterator(iterator):
@@ -253,6 +270,7 @@ class TensorFlowTrainer(base_trainer.Trainer):
             one_step_on_data = tf.function(
                 one_step_on_data, reduce_retracing=True, jit_compile=True
             )
+        one_step_on_data = self._autoconvert_optionals(one_step_on_data)
 
         @tf.autograph.experimental.do_not_convert
         def one_step_on_data_distributed(data):
@@ -409,7 +427,7 @@ class TensorFlowTrainer(base_trainer.Trainer):
                     _use_cached_eval_dataset=True,
                 )
                 val_logs = {
-                    "val_" + name: val for name, val in val_logs.items()
+                    f"val_{name}": val for name, val in val_logs.items()
                 }
                 epoch_logs.update(val_logs)
 

keras/src/backend/torch/core.py

@@ -191,21 +191,18 @@ def convert_to_tensor(x, dtype=None, sparse=None, ragged=None):
         raise ValueError("`sparse=True` is not supported with torch backend")
     if ragged:
         raise ValueError("`ragged=True` is not supported with torch backend")
-    if isinstance(x, Variable):
-        if dtype is None:
-            return x.value
-        x = x.value
-        return x.to(to_torch_dtype(dtype))
-    if is_tensor(x):
+    if isinstance(x, Variable) or is_tensor(x):
+        if isinstance(x, Variable):
+            x = x.value
         device = get_device()
         if x.device != device:
             if x.is_meta:
                 x = torch.empty_like(x, device=device)
             else:
                 x = x.to(device)
-        if dtype is None:
-            return x
-        return x.to(to_torch_dtype(dtype))
+        if dtype is not None:
+            x = x.to(to_torch_dtype(dtype))
+        return x
     if dtype is None:
         if isinstance(x, bool):
             return torch.as_tensor(x, dtype=torch.bool, device=get_device())

keras/src/backend/torch/nn.py

@@ -9,7 +9,6 @@ from keras.src.backend.torch.core import cast
 from keras.src.backend.torch.core import convert_to_tensor
 from keras.src.backend.torch.core import get_device
 from keras.src.backend.torch.numpy import expand_dims
-from keras.src.backend.torch.numpy import maximum
 from keras.src.backend.torch.numpy import where
 from keras.src.utils.argument_validation import standardize_tuple
 
@@ -668,7 +667,7 @@ def one_hot(x, num_classes, axis=-1, dtype=None, sparse=False):
     # manual handling for negatives in the input to one_hot by using max(x, 0).
     # The output will have some invalid results, so we set them back to 0 using
     # `where` afterwards.
-    output = tnn.one_hot(maximum(x, 0), num_classes)
+    output = tnn.one_hot(torch.clamp(x, min=0), num_classes)
     output = where(expand_dims(x, axis=-1) >= 0, output, zero)
     output = convert_to_tensor(output, dtype=dtype)
     dims = output.dim()

keras/src/backend/torch/numpy.py

@@ -854,6 +854,22 @@ def hstack(xs):
     return torch.hstack(xs)
 
 
+def hypot(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+
+    dtype = dtypes.result_type(x1.dtype, x2.dtype)
+    if dtype in ["int8", "int16", "int32", "uint8", "uint16", "uint32"]:
+        dtype = config.floatx()
+    elif dtype == "int64":
+        dtype = "float64"
+
+    x1 = cast(x1, dtype)
+    x2 = cast(x2, dtype)
+
+    return torch.hypot(x1, x2)
+
+
 def identity(n, dtype=None):
     dtype = to_torch_dtype(dtype or config.floatx())
 

keras/src/backend/torch/trainer.py

@@ -299,7 +299,7 @@ class TorchTrainer(base_trainer.Trainer):
                     _use_cached_eval_dataset=True,
                 )
                 val_logs = {
-                    "val_" + name: val for name, val in val_logs.items()
+                    f"val_{name}": val for name, val in val_logs.items()
                 }
                 epoch_logs.update(val_logs)
 

keras/src/callbacks/backup_and_restore.py

@@ -99,9 +99,9 @@ class BackupAndRestore(Callback):
         self._training_metadata_path = file_utils.join(
             backup_dir, "training_metadata.json"
         )
-        self._prev_weights_path = self._weights_path + ".bkp"
+        self._prev_weights_path = f"{self._weights_path}.bkp"
         self._prev_training_metadata_path = (
-            self._training_metadata_path + ".bkp"
+            f"{self._training_metadata_path}.bkp"
         )
         if save_freq != "epoch" and not isinstance(save_freq, int):
             raise ValueError(

keras/src/callbacks/csv_logger.py

@@ -79,7 +79,7 @@ class CSVLogger(Callback):
                     val_keys_found = True
                     break
             if not val_keys_found and self.keys:
-                self.keys.extend(["val_" + k for k in self.keys])
+                self.keys.extend([f"val_{k}" for k in self.keys])
 
         if not self.writer:
 

keras/src/callbacks/model_checkpoint.py

@@ -372,7 +372,7 @@ class ModelCheckpoint(MonitorCallback):
         """
         dir_name = os.path.dirname(pattern)
         base_name = os.path.basename(pattern)
-        base_name_regex = "^" + re.sub(r"{.*}", r".*", base_name) + "$"
+        base_name_regex = f"^{re.sub(r'{.*}', r'.*', base_name)}$"
 
         latest_mod_time = 0
         file_path_with_latest_mod_time = None

keras/src/callbacks/tensorboard.py

@@ -424,7 +424,7 @@ class TensorBoard(Callback):
             with self._val_writer.as_default():
                 for name, value in logs.items():
                     self.summary.scalar(
-                        "evaluation_" + name + "_vs_iterations",
+                        f"evaluation_{name}_vs_iterations",
                         value,
                         step=self.model.optimizer.iterations,
                     )
@@ -460,7 +460,7 @@ class TensorBoard(Callback):
         if isinstance(logs, dict):
             for name, value in logs.items():
                 self.summary.scalar(
-                    "batch_" + name, value, step=self._global_train_batch
+                    f"batch_{name}", value, step=self._global_train_batch
                 )
 
         if not self._should_trace:
@@ -548,12 +548,12 @@ class TensorBoard(Callback):
         if train_logs:
             with self._train_writer.as_default():
                 for name, value in train_logs.items():
-                    self.summary.scalar("epoch_" + name, value, step=epoch)
+                    self.summary.scalar(f"epoch_{name}", value, step=epoch)
         if val_logs:
             with self._val_writer.as_default():
                 for name, value in val_logs.items():
                     name = name[4:]  # Remove 'val_' prefix.
-                    self.summary.scalar("epoch_" + name, value, step=epoch)
+                    self.summary.scalar(f"epoch_{name}", value, step=epoch)
 
     def _log_weights(self, epoch):
         """Logs the weights of the Model to TensorBoard."""
@@ -562,14 +562,14 @@ class TensorBoard(Callback):
                 for weight in layer.weights:
                     weight_name = weight.name.replace(":", "_")
                     # Add a suffix to prevent summary tag name collision.
-                    histogram_weight_name = weight_name + "/histogram"
+                    histogram_weight_name = f"{weight_name}/histogram"
                     self.summary.histogram(
                         histogram_weight_name, weight, step=epoch
                     )
                     if self.write_images:
                         # Add a suffix to prevent summary tag name
                         # collision.
-                        image_weight_name = weight_name + "/image"
+                        image_weight_name = f"{weight_name}/image"
                         self._log_weight_as_image(
                             weight, image_weight_name, epoch
                         )

keras/src/constraints/constraints.py

@@ -110,7 +110,9 @@ class MaxNorm(Constraint):
         w = backend.convert_to_tensor(w)
         norms = ops.sqrt(ops.sum(ops.square(w), axis=self.axis, keepdims=True))
         desired = ops.clip(norms, 0, self.max_value)
-        return w * (desired / (backend.epsilon() + norms))
+        return ops.cast(w, norms.dtype) * (
+            desired / (backend.epsilon() + norms)
+        )
 
     def get_config(self):
         return {"max_value": self.max_value, "axis": self.axis}
@@ -122,7 +124,7 @@ class NonNeg(Constraint):
 
     def __call__(self, w):
         w = backend.convert_to_tensor(w)
-        return w * ops.cast(ops.greater_equal(w, 0.0), dtype=w.dtype)
+        return ops.multiply(w, ops.greater_equal(w, 0.0))
 
 
 @keras_export(["keras.constraints.UnitNorm", "keras.constraints.unit_norm"])
@@ -148,10 +150,8 @@ class UnitNorm(Constraint):
 
     def __call__(self, w):
         w = backend.convert_to_tensor(w)
-        return w / (
-            backend.epsilon()
-            + ops.sqrt(ops.sum(ops.square(w), axis=self.axis, keepdims=True))
-        )
+        norms = ops.sqrt(ops.sum(ops.square(w), axis=self.axis, keepdims=True))
+        return ops.cast(w, norms.dtype) / (backend.epsilon() + norms)
 
     def get_config(self):
         return {"axis": self.axis}
@@ -202,7 +202,9 @@ class MinMaxNorm(Constraint):
             self.rate * ops.clip(norms, self.min_value, self.max_value)
             + (1 - self.rate) * norms
         )
-        return w * (desired / (backend.epsilon() + norms))
+        return ops.cast(w, norms.dtype) * (
+            desired / (backend.epsilon() + norms)
+        )
 
     def get_config(self):
         return {

keras/src/datasets/boston_housing.py

@@ -48,7 +48,7 @@ def load_data(path="boston_housing.npz", test_split=0.2, seed=113):
     )
     path = get_file(
         path,
-        origin=origin_folder + "boston_housing.npz",
+        origin=f"{origin_folder}boston_housing.npz",
         file_hash=(  # noqa: E501
             "f553886a1f8d56431e820c5b82552d9d95cfcb96d1e678153f8839538947dff5"
         ),

keras/src/datasets/california_housing.py

@@ -73,7 +73,7 @@ def load_data(
     )
     path = get_file(
         path,
-        origin=origin_folder + "california_housing.npz",
+        origin=f"{origin_folder}california_housing.npz",
         file_hash=(  # noqa: E501
             "1a2e3a52e0398de6463aebe6f4a8da34fb21fbb6b934cf88c3425e766f2a1a6f"
         ),

keras/src/datasets/cifar10.py

@@ -79,7 +79,7 @@ def load_data():
     # batches are within an inner folder
     path = os.path.join(path, "cifar-10-batches-py")
     for i in range(1, 6):
-        fpath = os.path.join(path, "data_batch_" + str(i))
+        fpath = os.path.join(path, f"data_batch_{i}")
         (
             x_train[(i - 1) * 10000 : i * 10000, :, :, :],
             y_train[(i - 1) * 10000 : i * 10000],

keras/src/datasets/cifar100.py

@@ -71,10 +71,10 @@ def load_data(label_mode="fine"):
 
     path = os.path.join(path, "cifar-100-python")
     fpath = os.path.join(path, "train")
-    x_train, y_train = load_batch(fpath, label_key=label_mode + "_labels")
+    x_train, y_train = load_batch(fpath, label_key=f"{label_mode}_labels")
 
     fpath = os.path.join(path, "test")
-    x_test, y_test = load_batch(fpath, label_key=label_mode + "_labels")
+    x_test, y_test = load_batch(fpath, label_key=f"{label_mode}_labels")
 
     y_train = np.reshape(y_train, (len(y_train), 1))
     y_test = np.reshape(y_test, (len(y_test), 1))

keras/src/datasets/imdb.py

@@ -78,7 +78,7 @@ def load_data(
     )
     path = get_file(
         fname=path,
-        origin=origin_folder + "imdb.npz",
+        origin=f"{origin_folder}imdb.npz",
         file_hash=(  # noqa: E501
             "69664113be75683a8fe16e3ed0ab59fda8886cb3cd7ada244f7d9544e4676b9f"
         ),
@@ -181,7 +181,7 @@ def get_word_index(path="imdb_word_index.json"):
     )
     path = get_file(
         fname=path,
-        origin=origin_folder + "imdb_word_index.json",
+        origin=f"{origin_folder}imdb_word_index.json",
         file_hash="bfafd718b763782e994055a2d397834f",
     )
     with open(path) as f:

keras/src/datasets/mnist.py

@@ -59,7 +59,7 @@ def load_data(path="mnist.npz"):
     )
     path = get_file(
         fname=path,
-        origin=origin_folder + "mnist.npz",
+        origin=f"{origin_folder}mnist.npz",
         file_hash=(  # noqa: E501
             "731c5ac602752760c8e48fbffcf8c3b850d9dc2a2aedcf2cc48468fc17b673d1"
         ),

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/initializers/constant_initializers.py

@@ -253,14 +253,18 @@ class STFT(Initializer):
                 scaling = ops.sum(ops.abs(win))
 
         _fft_length = (fft_length - 1) * 2
-        freq = (
-            ops.reshape(ops.arange(fft_length, dtype=dtype), (1, 1, fft_length))
-            / _fft_length
+        freq = ops.divide(
+            ops.reshape(
+                ops.arange(fft_length, dtype=dtype), (1, 1, fft_length)
+            ),
+            _fft_length,
         )
         time = ops.reshape(
             ops.arange(frame_length, dtype=dtype), (frame_length, 1, 1)
         )
-        args = -2 * time * freq * ops.arccos(ops.cast(-1, dtype))
+        args = ops.multiply(ops.multiply(-2, time), freq) * ops.arccos(
+            ops.cast(-1, dtype)
+        )
 
         if self.side == "real":
             kernel = ops.cast(ops.cos(args), dtype)
@@ -268,7 +272,7 @@ class STFT(Initializer):
             kernel = ops.cast(ops.sin(args), dtype)
 
         if win is not None:
-            kernel = kernel * win / scaling
+            kernel = ops.divide(ops.multiply(kernel, win), scaling)
         return kernel
 
     def get_config(self):

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: