Grappler ImplementationSelector fails to use CUDNN (for GRU & LSTM) when using tf.data.experimental.scan

[OmriSteiner]

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): Ubuntu 20.04
• TensorFlow installed from (source or binary): binary (latest-gpu docker)
• TensorFlow version (use command below): v2.5.0-rc3-213-ga4dfb8d1a71 2.5.0 (but this also reproduces in earlier versions).
• Python version: 3.6.9
• CUDA/cuDNN version: 11.2 / 8.1.0
• GPU model and memory: NVIDIA Quadro P1000

Describe the current behavior
As you can see in the code here, the way CUDNN LSTM & GRU are implemented, is by adding two functions with the same name to the graph. Grappler's ImplementationSelector optimizer runs later and selects the correct implementation. However, when using tf.data.experimental.scan, from what I've previously seen (I don't quite remember well), the function gets renamed - which causes the implementation selector to fail.
What actually happens is the non-CUDNN version is the default implementation - so it falls back to that one.

Describe the expected behavior
The layers should run using CUDNN, when possible.

Contributing

• Do you want to contribute a PR? (yes/no): This might be a bit over my head, but if the solution is simple, I could implement it.
• Briefly describe your candidate solution(if contributing):

Standalone code to reproduce the issue
Unfortunately I couldn't think of a way to display this issue without Tensorboard, but here it is:

import tensorflow as tf

model = tf.keras.models.Sequential([
    tf.keras.layers.GRU(4),
])

@tf.function
def foo(dataset):
    a = 1.
    # Autograph converts this to tf.data.experimental.scan
    for i in dataset:
        a = model(i)
    return a # return the value to avoid pruning

inp_tensor = tf.random.normal([32, 10, 8])
dataset = tf.data.Dataset.from_tensor_slices([inp_tensor])

tf.profiler.experimental.start("/tmp/tensorboard")

# This successfully uses CUDNN
model.predict(inp_tensor)

# This doesn't use CUDNN
foo(dataset)

tf.profiler.experimental.stop()

Here you can see model.predict uses CUDNN:

And here you can see foo(dataset) does not use CUDNN:

[tilakrayal]

I was able to reproduce the code in tf v2.4 and v2.5.Please find the gist of it here.