Keras fails to compute gradients for autoencoder-esce model using Tensorflow backend with mixed precision and jit compilation enabled. See code here: colab.
This is caused by UpSampling2D layer. When gradients are computed, the type is resolved as float32 instead of float16, and this causes Relu that comes next to throw a dtype mismatch exception.
The only working workaround I found is explicitly setting dtype to float32 for UpSampling2D layer. This inserts a cast node inbetween relu and upsample which helps dealing with dtype conversion.
Not sure which project this issue should be submitted to: Keras, TF or XLA
Comment From: divyashreepathihalli
Looks like tf.image.resize(..., method='bilinear') gradient op produces float 32 tensors
Keras casts forward pass output to float16
consequent layers produce float16 gradients - leading to XLA error
Fix would need to re-implement tf.image.resize(..., method='bilinear`)
Works fine on other backends
This is a tensorflow issue. A temporaray worksaaround it to set the dtype to "float32" on that layer
Comment From: github-actions[bot]
This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.
Comment From: itmo153277
Ping to avoid stale status