Keras fails to load saved model / properly infer dtypes in tf.math.maximum

[mikolajpabiszczak]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS Big Sur 11.1
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N/A
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): v2.4.0-rc4-71-g582c8d236cb 2.4.0
• Python version: 3.8.6
• Bazel version (if compiling from source): N/A
• GCC/Compiler version (if compiling from source): N/A
• CUDA/cuDNN version: N/A
• GPU model and memory: N/A

Describe the current behavior
Keras fails to load model due to problems with inferring data types in tf.math.maximum. In particular: take the input to the network to be float32, then cast the tensor to float64 and feed into maximum layer (with the second input to that layer being float64 constant tensor) - creating this model is successful. Save that model (below I used SavedModel format, but .h5 case is similar) and try to load it - error is raised.

Moreover, the error is raised only when the constant tensor is passed as first input to tf.math.maxiumum and does not occur if it is passed as the second input - see the Colab notebook I attached below.

I believe this is strictly Keras-related, as I was able to successfully convert this failing-to-load model to to .tflite version (with proper options of converter set, using from_saved_model method) and this tflite model works and its' data types are correct.

Describe the expected behavior
The model should load properly.

Standalone code to reproduce the issue

This code snipper should reproduce the issue:

import tensorflow as tf
from tensorflow import keras

inp = keras.Input(shape=(1))
x = tf.cast(inp, dtype=tf.float64)
a = tf.constant(1.0, dtype=tf.float64)
x = tf.maximum(a, x)
out = tf.cast(x, dtype=tf.float32)

model = keras.models.Model(inp, out)
model.summary()

model.save('dummy_model')
del model

loaded_model = tf.keras.models.load_model('dummy_model')

BTW: When using x = tf.maximum(x, a) instead of x = tf.maximum(a, x) in the example above, the error is not raised!

See also slightly more elaborate Colab notebook.

Other info / logs

2021-02-15 12:45:01.772147: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-02-15 12:45:01.772422: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
Model: "model"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 1)]               0         
_________________________________________________________________
tf.cast (TFOpLambda)         (None, 1)                 0         
_________________________________________________________________
tf.math.maximum (TFOpLambda) (None, 1)                 0         
_________________________________________________________________
tf.cast_1 (TFOpLambda)       (None, 1)                 0         
=================================================================
Total params: 0
Trainable params: 0
Non-trainable params: 0
_________________________________________________________________
2021-02-15 12:45:01.835730: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
Traceback (most recent call last):
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/framework/op_def_library.py", line 517, in _apply_op_helper
    values = ops.convert_to_tensor(
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/profiler/trace.py", line 163, in wrapped
    return func(*args, **kwargs)
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 1507, in convert_to_tensor
    raise ValueError(
ValueError: Tensor conversion requested dtype float32 for Tensor with dtype float64: <tf.Tensor 'Placeholder:0' shape=(None, 1) dtype=float64>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "tmp.py", line 32, in <module>
    loaded_model = tf.keras.models.load_model('dummy_model')
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/saving/save.py", line 212, in load_model
    return saved_model_load.load(filepath, compile, options)
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 147, in load
    keras_loader.finalize_objects()
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 600, in finalize_objects
    self._reconstruct_all_models()
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 619, in _reconstruct_all_models
    self._reconstruct_model(model_id, model, layers)
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 665, in _reconstruct_model
    created_layers) = functional_lib.reconstruct_from_config(
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/engine/functional.py", line 1285, in reconstruct_from_config
    process_node(layer, node_data)
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/engine/functional.py", line 1233, in process_node
    output_tensors = layer(input_tensors, **kwargs)
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 951, in __call__
    return self._functional_construction_call(inputs, args, kwargs,
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 1090, in _functional_construction_call
    outputs = self._keras_tensor_symbolic_call(
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 822, in _keras_tensor_symbolic_call
    return self._infer_output_signature(inputs, args, kwargs, input_masks)
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 863, in _infer_output_signature
    outputs = call_fn(inputs, *args, **kwargs)
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/layers/core.py", line 1327, in _call_wrapper
    return self._call_wrapper(*args, **kwargs)
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/keras/layers/core.py", line 1359, in _call_wrapper
    result = self.function(*args, **kwargs)
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py", line 5704, in maximum
    _, _, _op, _outputs = _op_def_library._apply_op_helper(
  File "/Users/mikolajpabiszczak/.pyenv/versions/abnormal-sounds-lite-p3.8.6-tf2.4/lib/python3.8/site-packages/tensorflow/python/framework/op_def_library.py", line 555, in _apply_op_helper
    raise TypeError(
TypeError: Input 'y' of 'Maximum' Op has type float64 that does not match type float32 of argument 'x'.

[amahendrakar]

Facing an error stating ValueError: 2 errors while building NodeDef 'tf_op_layer_Maximum/Maximum' using Op<name=Maximum; signature=x:T, y:T -> z:T; attr=T:type,allowed=[DT_BFLOAT16, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_UINT8, DT_INT16, DT_INT32, DT_INT64]>: on running the code with TF v2.3.

Was able to reproduce the issue with TF v2.4 and TF-nightly.

Please check the linked gist for reference. Thanks!

[sushreebarsa]

I was able to reproduce the issue with tf v2.5 ,Please check the gist ..Thanks !

[mikolajpabiszczak]

Just bumping here: it still persists.

[tilakrayal]

I was able to reproduce the issue on tf-nightly 2.10.0-dev20220719. Kindly find the gist of it here. Thank you!