Repository: keras-team/keras-core

Issue:
Torch GPU test failures
Install PyTorch via - 
```
conda install pytorch torchvision pytorch-cuda=11.8 -c pytorch -c nvidia
```

We still have to install TF (for tf.data) and JAX (for import - though this could be avoided) in the environment via - 
```
pip install tensorflow_cpu
pip install jax[cpu]
```

Test via -
```
 KERAS_BACKEND=torch pytest keras_core --ignore keras_core/applications
```
Result summary:
```
FAILED keras_core/layers/convolutional/conv_test.py::ConvCorrectnessTest::test_conv3d0 - AssertionError:
FAILED keras_core/layers/convolutional/conv_test.py::ConvCorrectnessTest::test_conv3d1 - AssertionError:
FAILED keras_core/layers/convolutional/conv_transpose_test.py::ConvTransposeCorrectnessTest::test_conv3d_transpose2 - AssertionError:
FAILED keras_core/layers/preprocessing/random_brightness_test.py::RandomBrightnessTest::test_tf_data_compatibility - TypeError: in user code:
FAILED keras_core/layers/preprocessing/random_contrast_test.py::RandomContrastTest::test_tf_data_compatibility - TypeError: in user code:
```

ConvCorrectnessTest, ConvTransposeCorrectnessTest - Failed due to mismatch in test array values.
Preprpcoessing RandomBrightnessTest, RandomContrastTest failed on `numpy()` being calld on GPU tensor.  It needs to be `cpu().numpy()` for torch.


Discussion:
Failure details for converting the tensor to `cpu` in Random Brightness and Contrast -
```
E               File "code/keras-core/keras_core/layers/preprocessing/random_brightness.py", line 121, in call  *
E                   return self._brightness_adjust(inputs)
E               File "code/keras-core/keras_core/layers/preprocessing/random_brightness.py", line 144, in _brightness_adjust  *
E                   rgb_delta = self.backend.random.uniform(
E               File "code/keras-core/keras_core/backend/tensorflow/random.py", line 25, in uniform  *
E                   seed = tf_draw_seed(seed)
E               File "code/keras-core/keras_core/backend/tensorflow/random.py", line 12, in tf_draw_seed  *
E                   return tf.cast(draw_seed(seed), dtype="int32")
E               File "lib/python3.10/site-packages/torch/_tensor.py", line 970, in __array__
E                   return self.numpy()
E
E               TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

lib/python3.10/site-packages/torch/_tensor.py:970: TypeError
___________________________________________ RandomContrastTest.test_tf_data_compatibility
```
Thanks for the report.

> ConvCorrectnessTest, ConvTransposeCorrectnessTest - Failed due to mismatch in test array values.


What's the magnitude of the difference? It's expected that GPU kernels outputs should differ from CPU outputs by a little, typically on the order of 1e-7.
> Failure details for converting the tensor to cpu in Random Brightness and Contrast 

The failed tests are for `tf.data` compat specifically. It looks like we try to use a SeedGenerator created in torch to feed the TF random seed. Even if it did work it would be incorrect. We'd need to use a TF-created SeedGenerator.

We already have a fix for this in the code (look at the 4 lines above line 144 where the failure occurs) so I'm curious why the fix isn't operating here.


Yes, those differences are mostly small.  But sometimes big (like second value in array) - -6.273717e-03 vs -6.138057e-03
```
________________________________________ ConvTransposeCorrectnessTest.test_conv3d_transpose2 ________________________________________

self = <keras_core.layers.convolutional.conv_transpose_test.ConvTransposeCorrectnessTest testMethod=test_conv3d_transpose2>
filters = 6, kernel_size = (2, 2, 3), strides = (2, 1, 2), padding = 'valid', output_padding = None, data_format = 'channels_last'
dilation_rate = (1, 1, 1)

    @parameterized.parameters(
        {
            "filters": 5,
            "kernel_size": 2,
            "strides": 2,
            "padding": "valid",
            "output_padding": None,
            "data_format": "channels_last",
            "dilation_rate": 1,
        },
        {
            "filters": 6,
            "kernel_size": 2,
            "strides": 3,
            "padding": "same",
            "output_padding": 2,
            "data_format": "channels_last",
            "dilation_rate": (1, 1, 1),
        },
        {
            "filters": 6,
            "kernel_size": (2, 2, 3),
            "strides": (2, 1, 2),
            "padding": "valid",
            "output_padding": None,
            "data_format": "channels_last",
            "dilation_rate": (1, 1, 1),
        },
    )
    def test_conv3d_transpose(
        self,
        filters,
        kernel_size,
        strides,
        padding,
        output_padding,
        data_format,
        dilation_rate,
    ):
        layer = layers.Conv3DTranspose(
            filters=filters,
            kernel_size=kernel_size,
            strides=strides,
            padding=padding,
            output_padding=output_padding,
            data_format=data_format,
            dilation_rate=dilation_rate,
        )
        tf_keras_layer = tf.keras.layers.Conv3DTranspose(
            filters=filters,
            kernel_size=kernel_size,
            strides=strides,
            padding=padding,
            output_padding=output_padding,
            data_format=data_format,
            dilation_rate=dilation_rate,
        )

        inputs = np.random.normal(size=[2, 8, 8, 8, 4])
        layer.build(input_shape=inputs.shape)
        tf_keras_layer.build(input_shape=inputs.shape)

        kernel_shape = layer.kernel.shape
        kernel_weights = np.random.normal(size=kernel_shape)
        bias_weights = np.random.normal(size=(filters,))
        layer.kernel.assign(kernel_weights)
        tf_keras_layer.kernel.assign(kernel_weights)

        layer.bias.assign(bias_weights)
        tf_keras_layer.bias.assign(bias_weights)
        outputs = layer(inputs)
        expected = tf_keras_layer(inputs)
>       self.assertAllClose(outputs, expected, atol=1e-5)

keras_core/layers/convolutional/conv_transpose_test.py:453:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
keras_core/testing/test_case.py:34: in assertAllClose
    np.testing.assert_allclose(x1, x2, atol=atol, rtol=rtol)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

args = (<function assert_allclose.<locals>.compare at 0x7f44d43979a0>, array([[[[[ 1.80561197e+00, -6.27371669e-03, -1.049309...94704497e+00, -5.35111713e+00,
            2.91078305e+00, -2.94839787e+00,  4.89157629e+00]]]]],
      dtype=float32))
kwds = {'equal_nan': True, 'err_msg': '', 'header': 'Not equal to tolerance rtol=1e-06, atol=1e-05', 'verbose': True}

    @wraps(func)
    def inner(*args, **kwds):
        with self._recreate_cm():
>           return func(*args, **kwds)
E           AssertionError:
E           Not equal to tolerance rtol=1e-06, atol=1e-05
E
E           Mismatched elements: 28944 / 29376 (98.5%)
E           Max absolute difference: 0.00715119
E           Max relative difference: 35.82705
E            x: array([[[[[ 1.805612e+00, -6.273717e-03, -1.049310e+00,  4.257508e+00,
E                      -5.621735e-01,  4.334869e-01],
E                     [-1.109822e+00, -2.608716e+00, -3.800041e+00,  4.270791e+00,...
E            y: array([[[[[ 1.805492e+00, -6.138057e-03, -1.049304e+00,  4.256462e+00,
E                      -5.625696e-01,  4.334863e-01],
E                     [-1.109615e+00, -2.608501e+00, -3.800791e+00,  4.270457e+00,...
```
> The failed tests are for `tf.data` compat specifically. It looks like we try to use a SeedGenerator created in torch to feed the TF random seed. Even if it did work it would be incorrect. We'd need to use a TF-created SeedGenerator.

The issue looks to be although we might be getting seed generator from TF backend `keras_core.backend.tensorflow.random`, it redirects to `keras_core.random.SeedGenerator` that defaults back to environment backend (https://github.com/keras-team/keras-core/blob/main/keras_core/backend/tensorflow/random.py#L5).
Hence the state variable and increments are Torch (or JAX) tensors. Do we need to change the SeedGenerator to also take an optional Backend so that we can override the backend?  Currently SeedGenerator class uses `keras_core.backend` when creating Variable, which is the backend of the environment.
This should be fixed now. Please check. Thanks for the analysis, it enabled me to fix the issue quickly.

Resolving patch:
diff --git a/keras_core/layers/preprocessing/tf_data_layer.py b/keras_core/layers/preprocessing/tf_data_layer.py
--- a/keras_core/layers/preprocessing/tf_data_layer.py
+++ b/keras_core/layers/preprocessing/tf_data_layer.py
@@ -1,6 +1,6 @@
 from tensorflow import nest
 
-from keras_core import backend
+import keras_core.backend
 from keras_core.layers.layer import Layer
 from keras_core.random.seed_generator import SeedGenerator
 from keras_core.utils import backend_utils
@@ -22,7 +22,7 @@ def __init__(self, **kwargs):
 
     def __call__(self, inputs, **kwargs):
         if backend_utils.in_tf_graph() and not isinstance(
-            inputs, backend.KerasTensor
+            inputs, keras_core.backend.KerasTensor
         ):
             # We're in a TF graph, e.g. a tf.data pipeline.
             self.backend.set_backend("tensorflow")
@@ -47,7 +47,7 @@ def __call__(self, inputs, **kwargs):
 
     @tracking.no_automatic_dependency_tracking
     def _get_seed_generator(self, backend=None):
-        if backend is None or backend == self.backend._backend:
+        if backend is None or backend == keras_core.backend.backend():
             return self.generator
         if not hasattr(self, "_backend_generators"):
             self._backend_generators = {}
diff --git a/keras_core/random/seed_generator.py b/keras_core/random/seed_generator.py
--- a/keras_core/random/seed_generator.py
+++ b/keras_core/random/seed_generator.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 
+import keras_core.backend
 from keras_core.api_export import keras_core_export
 
 
@@ -30,9 +31,9 @@ def __init__(self, seed, **kwargs):
         if kwargs:
             raise ValueError(f"Unrecognized keyword arguments: {kwargs}")
         if custom_backend is not None:
-            backend = custom_backend
+            self.backend = custom_backend
         else:
-            from keras_core import backend
+            self.backend = keras_core.backend
 
         if seed is None:
             seed = make_default_seed()
@@ -43,9 +44,9 @@ def __init__(self, seed, **kwargs):
 
         def seed_initializer(*args, **kwargs):
             dtype = kwargs.get("dtype", None)
-            return backend.convert_to_tensor([seed, 0], dtype=dtype)
+            return self.backend.convert_to_tensor([seed, 0], dtype=dtype)
 
-        self.state = backend.Variable(
+        self.state = self.backend.Variable(
             seed_initializer,
             shape=(2,),
             dtype="uint32",
@@ -65,7 +66,9 @@ def draw_seed(seed):
         seed_state = seed.state
         # Use * 1 to create a copy
         new_seed_value = seed_state.value * 1
-        increment = convert_to_tensor(np.array([0, 1]), dtype="uint32")
+        increment = seed.backend.convert_to_tensor(
+            np.array([0, 1]), dtype="uint32"
+        )
         seed.state.assign(seed_state + increment)
         return new_seed_value
     elif isinstance(seed, int):


Test patch: