conv2d op produces output without nan with nan inputs
@pzy2000 If the lambda function or dropout layer is causing the NaNs, can we consider removing them and potentially adding alternative functionalities elsewhere in the network ?  The pre-processing data also needs to be verified once. Please use the latest stable TF version 2.16 and let us know?
Thank you! > @pzy2000 If the lambda function or dropout layer is causing the NaNs, can we consider removing them and potentially adding alternative functionalities elsewhere in the network? The pre-processing data also needs to be verified once. Please use the latest stable TF version 2.16 and let us know? Thank you!

Thanks for your answer! The reason why we use lambda function or dropout layer that is causing the NaNs is because we hope to detect framework defects like this with this method, since in my comprehension, given a conv2d ops, if input is NaNs, then output should also be NaNs. If my theory is mistaken, you are welcome to point it out! I will try it in latest version in a few hrs. @pzy2000 Could you please let us know if you have tried it in the latest TF? I tried it but faced a different issue, so could you please share all the dependencies for replicating?
Thank you! I just tried on latest tf, and it produced the following issue, of which I suppose was caused by inconsistent tf version:


Traceback (most recent call last):
  File "/root/zmx/COMET-master/issue/issue_tf_conv2d.py", line 218, in <module>
    model = keras.models.load_model(
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/saving/saving_api.py", line 183, in load_model
    return legacy_h5_format.load_model_from_hdf5(filepath)
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/legacy/saving/legacy_h5_format.py", line 133, in load_model_from_hdf5
    model = saving_utils.model_from_config(
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/legacy/saving/saving_utils.py", line 85, in model_from_config
    return serialization.deserialize_keras_object(
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/legacy/saving/serialization.py", line 495, in deserialize_keras_object
    deserialized_obj = cls.from_config(
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/models/model.py", line 528, in from_config
    return functional_from_config(
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/models/functional.py", line 509, in functional_from_config
    process_layer(layer_data)
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/models/functional.py", line 489, in process_layer
    layer = saving_utils.model_from_config(
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/legacy/saving/saving_utils.py", line 85, in model_from_config
    return serialization.deserialize_keras_object(
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/legacy/saving/serialization.py", line 495, in deserialize_keras_object
    deserialized_obj = cls.from_config(
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/layers/core/lambda_layer.py", line 193, in from_config
    fn = python_utils.func_load(
  File "/root/miniconda3/envs/tf_latest/lib/python3.9/site-packages/keras/src/utils/python_utils.py", line 83, in func_load
    code = marshal.loads(raw_code)
ValueError: bad marshal data (unknown type code)
 @pzy2000 The error occurs when the keras.models.load_model function attempts to deserialize a custom Lambda layer configuration saved in the model. It might be occurring for keras model loading. Please refer Keras documentation on saving and loading models [here](https://www.tensorflow.org/tutorials/keras/save_and_load ). 
Thank you! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/63391">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/63391">No</a>

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

binary

### TensorFlow version

tf 2.10.0

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 20.04.4

### Mobile device

_No response_

### Python version

3.7.6

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

12.2

### GPU model and memory

_No response_

### Current behavior?

![image](https://github.com/tensorflow/tensorflow/assets/77827720/d914a4d0-b420-4d29-b66d-06c1aa6ae1bb)
we expect layer 5 output to be also nan, since layer 5 input is nan
[lenet5-fashion-mnist_origin-SpecialI5-LMerg34.zip](https://github.com/tensorflow/tensorflow/files/14554559/lenet5-fashion-mnist_origin-SpecialI5-LMerg34.zip)


### Standalone code to reproduce the issue

```shell
import keras
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model


def _init_input(input_shape):
    input_shape = list(input_shape)
    input_shape[0] = 10
    input_shape = tuple(input_shape)
    input = np.random.rand(*input_shape)
    return input


class CustomCastLayer(keras.layers.Layer):
    def __init__(self, target_dtype, **kwargs):
        self.target_dtype = target_dtype
        super().__init__(**kwargs)

    def build(self, input_shape):
        """
        Do Nothing
        """
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        output = tf.cast(inputs, self.target_dtype)
        return output

    def get_config(self):
        config = super().get_config()
        config.update({"target_dtype": self.target_dtype})
        return config

    def compute_output_shape(self, input_shape):
        """
        Do Nothing
        """
        return input_shape


class CustomPadLayer(keras.layers.Layer):
    def __init__(self, padding=[[0, 0]], constant_values=0, **kwargs):
        self.padding = padding  # add [0,0] to padding so we will not change the shape of batch dimension
        self.constant_values = constant_values
        super().__init__(**kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        output = tf.pad(inputs, paddings=tf.constant([[0, 0]] + self.padding), mode="CONSTANT",
                        constant_values=self.constant_values)
        return output

    def get_config(self):
        config = super().get_config()
        config.update({"padding": self.padding, "constant_values": self.constant_values})
        return config

    def compute_output_shape(self, input_shape):
        """
        Formula to calculate the output shape
        Suppose the input_shape is (None, N, W, C), the `paddings` is [[a, b], [c, d]].
        The output_shape is (None, N+a+b, W+c+d, C).
        """
        input_shape_list = list(input_shape)
        padding = [[0, 0]] + self.padding
        assert len(input_shape_list) == len(padding)  # Two dimensions should match.
        output_shape = [None]
        for i, pad in zip(input_shape_list[1:], padding[1:]):
            output_shape.append(i + pad[0] + pad[1])
        return tuple(output_shape)


class CustomCropLayer(keras.layers.Layer):
    def __init__(self, cropping, **kwargs):
        self.cropping = cropping
        super().__init__(**kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        input_shape = inputs.shape.as_list()
        indices = [slice(None)]
        cropping = [[0, 0]] + self.cropping  # add [0,0] to padding so we will not change the shape of batch dimension
        for shape, crop in zip(input_shape[1:], cropping[1:]):
            indices.append(slice(0 + crop[0], shape - crop[1]))
        return inputs[indices]

    def get_config(self):
        config = super().get_config()
        config.update({"cropping": self.cropping})
        return config

    def compute_output_shape(self, input_shape):
        """
        Formula to calculate the output shape
        Suppose the input_shape is (None, N, W, C), the `cropping` is [[a, b], [c, d]].
        The output_shape is (None, N-a-b, W-c-d, C).
        """
        input_shape_list = list(input_shape)
        cropping = [[0, 0]] + self.cropping
        assert len(input_shape_list) == len(cropping)  # Two dimensions should match.
        output_shape = [None]
        for i, crop in zip(input_shape[1:], cropping[1:]):
            output_shape.append(i - crop[0] - crop[1])
        return tuple(output_shape)


class CustomExpandLayer(keras.layers.Layer):
    def __init__(self, axis=1, **kwargs):
        self.axis = axis
        super().__init__(**kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        return tf.expand_dims(inputs, self.axis)

    def get_config(self):
        config = super().get_config()
        config.update({"axis": self.axis})
        return config

    def compute_output_shape(self, input_shape):
        """
        Formula to calculate the output shape
        Suppose the input_shape is [None, N, W, C]:
        axis=0:
            output_shape: [1, None, N, W, C]
        axis=1 (default):
            output_shape: [None, 1, N, W, C]
        axis=2:
            output_shape: [None, N, 1, W, C]
        axis=3:
            output_shape: [None, N, W, 1, C]
        axis=4:
            output_shape: [None, N, W, C, 1]
        axis=5:
            raise Exception
        """
        input_shape_list = list(input_shape)
        if self.axis > len(input_shape_list):
            raise ValueError(f"axis {self.axis} should be smaller than input_shape + 1: {len(input_shape_list) + 1}")
        output_shape = input_shape_list[0:self.axis] + [1] + input_shape_list[self.axis:]
        return tuple(output_shape)  # we should use tuple!!! not list !!!


class CustomDropDimLayer(keras.layers.Layer):
    def __init__(self, axis=1, **kwargs):
        self.axis = axis
        super().__init__(**kwargs)

    def build(self, input_shape):
        super().build(input_shape)

    def call(self, inputs, **kwargs):
        """
        Something magic to automatically generate indices for array slicing.
        To determine a specific axis, we can use slice(None) to replace `:`
        """
        dim = len(inputs.shape)
        if self.axis > dim - 1 or self.axis < 1:
            raise ValueError(f"axis: {self.axis} should be within the range: [1, {dim - 1}] for {dim}D tensor")
        indices = [slice(None) for i in range(dim)]
        indices[self.axis] = 0
        return inputs[indices]

    def get_config(self):
        config = super().get_config()
        config.update({"axis": self.axis})
        return config

    def compute_output_shape(self, input_shape):
        """
        Formula to calculate the output shape
        Suppose the input_shape is [None, N, W, C]:
        axis=0:  # Although it is feasible, we don't allow this to happen
            Raise Exception
        axis=1 (default):
            output_shape: [None, W, C]
        axis=2:
            output_shape: [None, N, C]
        axis=3:
            output_shape: [None, N, W]
        axis=4:
            Raise Exception
        """
        input_shape_list = list(input_shape)
        output_shape = input_shape_list[0:self.axis] + input_shape_list[self.axis + 1:]
        return tuple(output_shape)


def custom_objects(mode="custom"):
    def no_activation(x):
        return x

    def leakyrelu(x):
        import keras.backend as K
        return K.relu(x, alpha=0.01)

    # objects = {}
    objects = {'no_activation': no_activation, 'leakyrelu': leakyrelu}
    if mode == "custom":
        objects['CustomPadLayer'] = CustomPadLayer
        objects['CustomCropLayer'] = CustomCropLayer
        objects['CustomDropDimLayer'] = CustomDropDimLayer
        objects['CustomExpandLayer'] = CustomExpandLayer
        objects['CustomCastLayer'] = CustomCastLayer

    return objects


# 加载预训练模型
model = keras.models.load_model(
    'lenet5-fashion-mnist_origin-SpecialI5-LMerg34.h5',
    custom_objects=custom_objects())

input_layer = model.input
layer_outputs = [layer.output for layer in model.layers]
print(layer_outputs)
print("===================================")

new_model = Model(inputs=input_layer, outputs=layer_outputs)
# 创建一个具有特定形状的随机tensor，例如形状为(1, 3, 32, 32)
with tf.device('cpu'):
    input_tensor = _init_input(model.input_shape)
    all_layer_outputs = new_model.predict(input_tensor)
    print("===================================")
    for i, layer_output in enumerate(all_layer_outputs):
        print(f"Name of layer {i}: {model.layers[i].name}")
        print(f"Output if nan of layer {i}: {np.isnan(layer_output).any()}")
        print()
```


### Relevant log output

```shell
Name of layer 0: conv2d_1_input
Output if nan of layer 0: False

Name of layer 1: conv2d_1_copy_SpecialI_copy_LMerg
Output if nan of layer 1: False

Name of layer 2: max_pooling2d_1_copy_SpecialI_copy_LMerg
Output if nan of layer 2: False

Name of layer 3: lambda_copy_LMerg
Output if nan of layer 3: True

Name of layer 4: dropout_1_copy_SpecialI_copy_LMerg
Output if nan of layer 4: True

Name of layer 5: conv2d_2_copy_SpecialI_copy_LMerg
Output if nan of layer 5: False

Name of layer 6: max_pooling2d_2_copy_SpecialI_copy_LMerg
Output if nan of layer 6: False

Name of layer 7: dropout_2_copy_SpecialI_copy_LMerg_merge1
Output if nan of layer 7: False

Name of layer 8: subtract_copy_ML
Output if nan of layer 8: False

Name of layer 9: flatten_1_copy_SpecialI_copy_LMerg
Output if nan of layer 9: False

Name of layer 10: dense_1_copy_SpecialI_copy_LMerg
Output if nan of layer 10: False

Name of layer 11: dense_2_copy_SpecialI_copy_LMerg
Output if nan of layer 11: False

Name of layer 12: dense_3_copy_SpecialI_copy_LMerg
Output if nan of layer 12: False

Name of layer 13: dense_insert
Output if nan of layer 13: False
```
