System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Colab
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.8
• Python version: 3.7

Describe the current behavior

When some loss (tf.losses.SparseCategoricalCrossentropy, tf.losses.CategoricalCrossentropy, tf.losses.BinaryCrossentropy, or tf.losses.MeanSquaredError) is used on Ragged tensors, which is computed via a tf.map_fn on a RaggedTensor, that the gradient computation on a GPU crashes with

Node: 'Adam/gradients/zeros_like_2'
2 root error(s) found.
  (0) INTERNAL:  No unary variant unary_op function found for op ZEROS_LIKE Variant type_name: RaggedTensorVariant for device type: GPU
	 [[{{node Adam/gradients/zeros_like_2}}]]
	 [[binary_crossentropy/map/while/loop_body_control/_124/_67]]
  (1) INTERNAL:  No unary variant unary_op function found for op ZEROS_LIKE Variant type_name: RaggedTensorVariant for device type: GPU
	 [[{{node Adam/gradients/zeros_like_2}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_16690]

The computation does not crash on a CPU and it does not crash when tf.functions are executed eagerly.

Also, if the tf.map_fn is circumvented by using the following argument to compile

  loss=lambda yt, yp: tf.losses.BinaryCrossentropy()(yt. values, yp.values)

it works on GPU without a crash.

Describe the expected behavior

The code does not crash on a GPU.

• Do you want to contribute a PR? (yes/no): no

Standalone code to reproduce the issue

A simple Colab reproducing the error is here: https://colab.research.google.com/drive/1OELAhvpQHhaz3sOYabf4SdBqKlQCjNjs?usp=sharing

Other info / logs

The map_fn used is here: https://github.com/keras-team/keras/blob/2db5acf3e3c5904b014cb409d3c514bef44f9640/keras/losses.py#L1408