tensorflow: Missing GPU op for zeros_like for RaggedTensorVariant, error occurs when Ragged Tensor fed thru tf.map_fn

System information

Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes, included below OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.10 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): pip binary TensorFlow version (use command below): v1.12.1-49539-g18d8bcbe72b 2.5.0-dev20210123 Python version: ‘3.8.6 | packaged by conda-forge | (default, Nov 27 2020, 19:31:52) \n[GCC 9.3.0]’ Bazel version (if compiling from source): n/a GCC/Compiler version (if compiling from source): n/a CUDA/cuDNN version: 11.0 / 8 GPU model and memory: TITAN X (Pascal) computeCapability: 6.1

Describe the current behavior

I have a keras layer RescaleB that accepts a ragged tensor with shape [batch, (time), in_dim]. The layer calls map_fn to process each example in the batch separately, scaling the values along the inner dimension by a trainable gain vector. (The details of the operation aren’t critical, but the ragged tensor going into map_fn is.)

Using this layer fails with No unary variant unary_op function found for unary variant op enum: 1 Variant type_name: RaggedTensorVariant for device type: GPU on a node whose name ends with rescale_b/map/while/TensorArrayV2Write/TensorListSetItem_grad/zeros_like which suggests that the zeros_like operation isn’t defined for Ragged Tensors on GPU?

In this simple example, i also include RescaleA, which accomplishes the same task using tf.ragged.map_flat_values, although in my real use case I need map_fn. This is a simplified example.

Describe the expected behavior

I’d expect RescaleB and RescaleA to function identically.

Standalone code to reproduce the issue

https://colab.research.google.com/drive/1mHycCXJL94VuCGkXIJ0bIXtbYamyZo78

I’ve reproduced the issue locally with tf-nightly-gpu TF 2.5, but I can’t seem to get the nightly version to see the GPU on Colab. The Colab notebook is using TF 2.4, but the issue remains in TF 2.5 nightly.

Other info / logs Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.

This may be the same issue as #44231 but hopefully the additional detail here is helpful.