Repository: keras-team/keras (python)

Issue:
keras with pytorch backend and mps set to default should use an mps generatir in randperm
Keras with pytorch backend and mps set to default needs to  use an mps generator in randperm

The following code

```
import os
os.environ["KERAS_BACKEND"] = "torch"

import torch as torch

torch.set_default_device('mps')

import keras
import numpy as np
from keras import layers

# Model / data parameters
num_classes = 10
input_shape = (28, 28, 1)

# Load the data and split it between train and test sets
(xx_train, yy_train), (xx_test, yy_test) = keras.datasets.mnist.load_data()

# Scale images to the [0, 1] range
x_train = xx_train.astype("float32") / 255
x_test = xx_test.astype("float32") / 255
# Make sure images have shape (28, 28, 1)
x_train = torch.from_numpy(np.expand_dims(xx_train, -1))
x_test = torch.from_numpy(np.expand_dims(xx_test, -1))
print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")


# convert class vectors to binary class matrices
y_train = torch.from_numpy(keras.utils.to_categorical(yy_train, num_classes).astype("float32"))
y_test = torch.from_numpy(keras.utils.to_categorical(yy_test, num_classes).astype("float32"))model = keras.Sequential(
    [
        keras.Input(shape=input_shape),
        layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        layers.MaxPooling2D(pool_size=(2, 2)),
        layers.Flatten(),
        layers.Dropout(0.5),
        layers.Dense(num_classes, activation="softmax"),
    ]
)
batch_size = 128
epochs = 15

model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)
```

produces the following error

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[5], line 6
      2 epochs = 15
      4 model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
----> 6 model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)

File [~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/keras/src/utils/traceback_utils.py:122](http://localhost:8888/~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/keras/src/utils/traceback_utils.py#line=121), in filter_traceback.<locals>.error_handler(*args, **kwargs)
    119     filtered_tb = _process_traceback_frames(e.__traceback__)
    120     # To get the full stack trace, call:
    121     # `keras.config.disable_traceback_filtering()`
--> 122     raise e.with_traceback(filtered_tb) from None
    123 finally:
    124     del filtered_tb

File [~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/torch/utils/data/dataloader.py:631](http://localhost:8888/~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/torch/utils/data/dataloader.py#line=630), in _BaseDataLoaderIter.__next__(self)
    628 if self._sampler_iter is None:
    629     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    630     self._reset()  # type: ignore[call-arg]
--> 631 data = self._next_data()
    632 self._num_yielded += 1
    633 if self._dataset_kind == _DatasetKind.Iterable and \
    634         self._IterableDataset_len_called is not None and \
    635         self._num_yielded > self._IterableDataset_len_called:

File ~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/torch/utils/data/dataloader.py:674, in _SingleProcessDataLoaderIter._next_data(self)
    673 def _next_data(self):
--> 674     index = self._next_index()  # may raise StopIteration
    675     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    676     if self._pin_memory:

File [~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/torch/utils/data/dataloader.py:621](http://localhost:8888/~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/torch/utils/data/dataloader.py#line=620), in _BaseDataLoaderIter._next_index(self)
    620 def _next_index(self):
--> 621     return next(self._sampler_iter)

File [~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/torch/utils/data/sampler.py:287](http://localhost:8888/~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/torch/utils/data/sampler.py#line=286), in BatchSampler.__iter__(self)
    285 batch = [0] * self.batch_size
    286 idx_in_batch = 0
--> 287 for idx in self.sampler:
    288     batch[idx_in_batch] = idx
    289     idx_in_batch += 1

File [~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/torch/utils/data/sampler.py:167](http://localhost:8888/~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/torch/utils/data/sampler.py#line=166), in RandomSampler.__iter__(self)
    165 else:
    166     for _ in range(self.num_samples // n):
--> 167         yield from torch.randperm(n, generator=generator).tolist()
    168     yield from torch.randperm(n, generator=generator).tolist()[:self.num_samples % n]

File ~/Pytorch/venv-Pytorch/lib/python3.12/site-packages/torch/utils/_device.py:77, in DeviceContext.__torch_function__(self, func, types, args, kwargs)
     75 if func in _device_constructors() and kwargs.get('device') is None:
     76     kwargs['device'] = self.device
---> 77 return func(*args, **kwargs)

RuntimeError: Expected a 'mps:0' generator device but found 'cpu'
```

Pull request:
Remove `PYTORCH_ENABLE_MPS_FALLBACK` flag requirement for mps
Running with MPS used to require this flag since the feature was experimental. It has since graduated to stable in PyTorch (see https://pytorch.org/docs/stable/notes/mps.html).

We can go ahead and remove the flag check. Fix #19436.

Patch:
diff --git a/keras/api/_tf_keras/keras/ops/__init__.py b/keras/api/_tf_keras/keras/ops/__init__.py
index 118c5a692..386730deb 100644
--- a/keras/api/_tf_keras/keras/ops/__init__.py
+++ b/keras/api/_tf_keras/keras/ops/__init__.py
@@ -72,6 +72,7 @@ from keras.src.ops.nn import moments
 from keras.src.ops.nn import multi_hot
 from keras.src.ops.nn import normalize
 from keras.src.ops.nn import one_hot
+from keras.src.ops.nn import psnr
 from keras.src.ops.nn import relu
 from keras.src.ops.nn import relu6
 from keras.src.ops.nn import selu
diff --git a/keras/api/_tf_keras/keras/ops/nn/__init__.py b/keras/api/_tf_keras/keras/ops/nn/__init__.py
index 61efc22a5..8c7e3d921 100644
--- a/keras/api/_tf_keras/keras/ops/nn/__init__.py
+++ b/keras/api/_tf_keras/keras/ops/nn/__init__.py
@@ -26,6 +26,7 @@ from keras.src.ops.nn import moments
 from keras.src.ops.nn import multi_hot
 from keras.src.ops.nn import normalize
 from keras.src.ops.nn import one_hot
+from keras.src.ops.nn import psnr
 from keras.src.ops.nn import relu
 from keras.src.ops.nn import relu6
 from keras.src.ops.nn import selu
diff --git a/keras/api/ops/__init__.py b/keras/api/ops/__init__.py
index 118c5a692..386730deb 100644
--- a/keras/api/ops/__init__.py
+++ b/keras/api/ops/__init__.py
@@ -72,6 +72,7 @@ from keras.src.ops.nn import moments
 from keras.src.ops.nn import multi_hot
 from keras.src.ops.nn import normalize
 from keras.src.ops.nn import one_hot
+from keras.src.ops.nn import psnr
 from keras.src.ops.nn import relu
 from keras.src.ops.nn import relu6
 from keras.src.ops.nn import selu
diff --git a/keras/api/ops/nn/__init__.py b/keras/api/ops/nn/__init__.py
index 61efc22a5..8c7e3d921 100644
--- a/keras/api/ops/nn/__init__.py
+++ b/keras/api/ops/nn/__init__.py
@@ -26,6 +26,7 @@ from keras.src.ops.nn import moments
 from keras.src.ops.nn import multi_hot
 from keras.src.ops.nn import normalize
 from keras.src.ops.nn import one_hot
+from keras.src.ops.nn import psnr
 from keras.src.ops.nn import relu
 from keras.src.ops.nn import relu6
 from keras.src.ops.nn import selu
diff --git a/keras/src/backend/jax/nn.py b/keras/src/backend/jax/nn.py
index 7fc623d83..3cbc61126 100644
--- a/keras/src/backend/jax/nn.py
+++ b/keras/src/backend/jax/nn.py
@@ -926,3 +926,16 @@ def ctc_decode(
             f"Invalid strategy {strategy}. Supported values are "
             "'greedy' and 'beam_search'."
         )
+
+
+def psnr(x1, x2, max_val):
+    if x1.shape != x2.shape:
+        raise ValueError(
+            f"Input shapes {x1.shape} and {x2.shape} must "
+            "match for PSNR calculation. "
+        )
+
+    max_val = convert_to_tensor(max_val, dtype=x2.dtype)
+    mse = jnp.mean(jnp.square(x1 - x2))
+    psnr = 20 * jnp.log10(max_val) - 10 * jnp.log10(mse)
+    return psnr
diff --git a/keras/src/backend/numpy/nn.py b/keras/src/backend/numpy/nn.py
index 2c27fad23..a2d89c323 100644
--- a/keras/src/backend/numpy/nn.py
+++ b/keras/src/backend/numpy/nn.py
@@ -967,3 +967,16 @@ def ctc_decode(
             f"Invalid strategy {strategy}. Supported values are "
             "'greedy' and 'beam_search'."
         )
+
+
+def psnr(x1, x2, max_val):
+    if x1.shape != x2.shape:
+        raise ValueError(
+            f"Input shapes {x1.shape} and {x2.shape} must "
+            "match for PSNR calculation. "
+        )
+
+    max_val = convert_to_tensor(max_val, dtype=x2.dtype)
+    mse = np.mean(np.square(x1 - x2))
+    psnr = 20 * np.log10(max_val) - 10 * np.log10(mse)
+    return psnr
diff --git a/keras/src/backend/tensorflow/nn.py b/keras/src/backend/tensorflow/nn.py
index 2a53c25f1..e7317599a 100644
--- a/keras/src/backend/tensorflow/nn.py
+++ b/keras/src/backend/tensorflow/nn.py
@@ -846,3 +846,18 @@ def ctc_decode(
     decoded_dense = tf.stack(decoded_dense, axis=0)
     decoded_dense = tf.cast(decoded_dense, "int32")
     return decoded_dense, scores
+
+
+def psnr(x1, x2, max_val):
+    from keras.src.backend.tensorflow.numpy import log10
+
+    if x1.shape != x2.shape:
+        raise ValueError(
+            f"Input shapes {x1.shape} and {x2.shape} must "
+            "match for PSNR calculation. "
+        )
+
+    max_val = convert_to_tensor(max_val, dtype=x2.dtype)
+    mse = tf.reduce_mean(tf.square(x1 - x2))
+    psnr = 20 * log10(max_val) - 10 * log10(mse)
+    return psnr
diff --git a/keras/src/backend/torch/core.py b/keras/src/backend/torch/core.py
index 257afeeec..8cc6c5b5b 100644
--- a/keras/src/backend/torch/core.py
+++ b/keras/src/backend/torch/core.py
@@ -1,5 +1,4 @@
 import contextlib
-import os
 
 import ml_dtypes
 import numpy as np
@@ -19,10 +18,7 @@ SUPPORTS_SPARSE_TENSORS = False
 # Some operators such as 'aten::_foreach_mul_.Scalar'
 # are not currently implemented for the MPS device.
 # check https://github.com/pytorch/pytorch/issues/77764.
-if (
-    torch.backends.mps.is_available()
-    and os.getenv("PYTORCH_ENABLE_MPS_FALLBACK") == "1"
-):
+if torch.backends.mps.is_available():
     DEFAULT_DEVICE = "mps"
 elif torch.cuda.is_available():
     DEFAULT_DEVICE = "cuda"
diff --git a/keras/src/backend/torch/nn.py b/keras/src/backend/torch/nn.py
index 4eef2c189..62749bc16 100644
--- a/keras/src/backend/torch/nn.py
+++ b/keras/src/backend/torch/nn.py
@@ -848,3 +848,20 @@ def ctc_decode(
             f"Invalid strategy {strategy}. Supported values are "
             "'greedy' and 'beam_search'."
         )
+
+
+def psnr(x1, x2, max_val):
+    if x1.shape != x2.shape:
+        raise ValueError(
+            f"Input shapes {x1.shape} and {x2.shape} must "
+            "match for PSNR calculation. "
+        )
+
+    x1, x2 = (
+        convert_to_tensor(x1),
+        convert_to_tensor(x2),
+    )
+    max_val = convert_to_tensor(max_val, dtype=x1.dtype)
+    mse = torch.mean((x1 - x2) ** 2)
+    psnr = 20 * torch.log10(max_val) - 10 * torch.log10(mse)
+    return psnr
diff --git a/keras/src/callbacks/early_stopping.py b/keras/src/callbacks/early_stopping.py
index e7c1fe9c0..5571cf606 100644
--- a/keras/src/callbacks/early_stopping.py
+++ b/keras/src/callbacks/early_stopping.py
@@ -50,7 +50,6 @@ class EarlyStopping(Callback):
             improvement is expected and thus training will not be stopped.
             Defaults to `0`.
 
-
     Example:
 
     >>> callback = keras.callbacks.EarlyStopping(monitor='loss',
diff --git a/keras/src/ops/nn.py b/keras/src/ops/nn.py
index b2b248154..1d900c8cc 100644
--- a/keras/src/ops/nn.py
+++ b/keras/src/ops/nn.py
@@ -2042,3 +2042,76 @@ def _normalize(x, axis=-1, order=2):
     norm = backend.linalg.norm(x, ord=order, axis=axis, keepdims=True)
     denom = backend.numpy.maximum(norm, epsilon)
     return backend.numpy.divide(x, denom)
+
+
+class PSNR(Operation):
+    def __init__(
+        self,
+        max_val,
+    ):
+        super().__init__()
+        self.max_val = max_val
+
+    def call(self, x1, x2):
+        return backend.nn.psnr(
+            x1=x1,
+            x2=x2,
+            max_val=self.max_val,
+        )
+
+    def compute_output_spec(self, x1, x2):
+        if len(x1.shape) != len(x2.shape):
+            raise ValueError("Inputs must have the same rank")
+
+        return KerasTensor(shape=())
+
+
+@keras_export(
+    [
+        "keras.ops.psnr",
+        "keras.ops.nn.psnr",
+    ]
+)
+def psnr(
+    x1,
+    x2,
+    max_val,
+):
+    """Peak Signal-to-Noise Ratio (PSNR) function.
+
+    This function computes the Peak Signal-to-Noise Ratio between two signals,
+    `x1` and `x2`. PSNR is a measure of the quality of a reconstructed signal.
+    The higher the PSNR, the closer the reconstructed signal is to the original
+    signal. Note that it can become negative when the signal power is
+    smaller that the noise power.
+
+    Args:
+        x1: The first input signal.
+        x2: The second input signal. Must have the same shape as `x1`.
+        max_val: The maximum possible value in the signals.
+
+    Returns:
+        float: The PSNR value between `x1` and `x2`.
+
+    Examples:
+
+    >>> x1 = keras.random.normal((2, 4, 4, 3))
+    >>> x2 = keras.random.normal((2, 4, 4, 3))
+    >>> max_val = 1.0
+    >>> keras.ops.nn.psnr(x1, x2, max_val)
+    -3.1697404
+    """
+    if any_symbolic_tensors(
+        (
+            x1,
+            x2,
+        )
+    ):
+        return PSNR(
+            max_val,
+        ).symbolic_call(x1, x2)
+    return backend.nn.psnr(
+        x1,
+        x2,
+        max_val,
+    )
diff --git a/requirements-common.txt b/requirements-common.txt
index 5d15f7ac6..f645c7ba9 100644
--- a/requirements-common.txt
+++ b/requirements-common.txt
@@ -1,4 +1,4 @@
-namex
+namex>=0.0.8
 black>=22
 flake8
 isort


Test patch:
diff --git a/keras/src/ops/nn_test.py b/keras/src/ops/nn_test.py
index 5e52a6c26..433bb0f46 100644
--- a/keras/src/ops/nn_test.py
+++ b/keras/src/ops/nn_test.py
@@ -654,6 +654,12 @@ class NNOpsDynamicShapeTest(testing.TestCase, parameterized.TestCase):
         x = KerasTensor([None, 2, 3])
         self.assertEqual(knn.normalize(x).shape, (None, 2, 3))
 
+    def test_psnr(self):
+        x1 = KerasTensor([None, 2, 3])
+        x2 = KerasTensor([None, 5, 6])
+        out = knn.psnr(x1, x2, max_val=224)
+        self.assertEqual(out.shape, ())
+
 
 class NNOpsStaticShapeTest(testing.TestCase):
     def test_relu(self):
@@ -1114,6 +1120,12 @@ class NNOpsStaticShapeTest(testing.TestCase):
         x = KerasTensor([1, 2, 3])
         self.assertEqual(knn.normalize(x).shape, (1, 2, 3))
 
+    def test_psnr(self):
+        x1 = KerasTensor([1, 2, 3])
+        x2 = KerasTensor([5, 6, 7])
+        out = knn.psnr(x1, x2, max_val=224)
+        self.assertEqual(out.shape, ())
+
 
 class NNOpsCorrectnessTest(testing.TestCase, parameterized.TestCase):
     def test_relu(self):
@@ -2032,6 +2044,25 @@ class NNOpsCorrectnessTest(testing.TestCase, parameterized.TestCase):
             ],
         )
 
+    def test_psnr(self):
+        x1 = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
+        x2 = np.array([[0.2, 0.2, 0.3], [0.4, 0.6, 0.6]])
+        max_val = 1.0
+        expected_psnr_1 = 20 * np.log10(max_val) - 10 * np.log10(
+            np.mean(np.square(x1 - x2))
+        )
+        psnr_1 = knn.psnr(x1, x2, max_val)
+        self.assertAlmostEqual(psnr_1, expected_psnr_1)
+
+        x3 = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
+        x4 = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
+        max_val = 1.0
+        expected_psnr_2 = 20 * np.log10(max_val) - 10 * np.log10(
+            np.mean(np.square(x3 - x4))
+        )
+        psnr_2 = knn.psnr(x3, x4, max_val)
+        self.assertAlmostEqual(psnr_2, expected_psnr_2)
+
 
 class NNOpsDtypeTest(testing.TestCase, parameterized.TestCase):
     """Test the dtype to verify that the behavior matches JAX."""