FFT produces wrong results when using multiple GPUs with MirroredStrategy
We are seeing similar issue recently with TF2.12, and it could be reproduced with this short script in a multi-gpu setup (we reproduced this on a machine with 2 A100 gpu, CUDA11.8, cuDNN8.6):

```
import numpy as np
import tensorflow as tf

# tf.debugging.set_log_device_placement(True)

print("multi gpu mirrored strategy test sanity check:")

def run_fft(img):
    print(tf.signal.fft2d(tf.experimental.numpy.copy(img)))

strategy = tf.distribute.MirroredStrategy()
rand_mat = tf.random.uniform(
                (1, 1, 686, 686),
                minval=0.0,
                maxval=1.0,
                dtype=tf.dtypes.float32,
                seed=None
            )
img = tf.cast(rand_mat, tf.complex64)
with strategy.scope():
    strategy.run(run_fft, args=(img,))
```
Output with TF2.9:
```
tf.Tensor(
[[[[ 2.35175906e+05  +0.j       -1.30122070e+02 +71.79061j
     2.00236359e+02 -40.830986j ...  2.86872749e+01-174.89594j
     2.00236313e+02 +40.831j    -1.30122101e+02 -71.790565j]
   [ 2.59338074e+01 -47.23604j   1.37966949e+02 +27.436068j
     1.20167496e+02+118.6681j   ...  1.09129379e+02 -69.504425j
     9.90038986e+01+118.647736j -6.46898956e+01 -98.40455j ]
   [ 1.74128876e+02+181.25378j  -2.27663452e+02 -70.84499j
    -3.25569305e+01+143.9032j   ...  1.96747189e+01-331.84482j
     8.72258377e+01+154.79916j   7.57820816e+01 -84.29112j ]
   ...
   [-8.93064423e+01 -94.69147j   2.13672241e+02+246.61185j
     2.80917282e+01-143.79309j  ...  2.39895905e+02 +28.519115j
     9.76210938e+01 -32.177147j  5.52456627e+01-292.4319j  ]
   [ 1.74128815e+02-181.2538j    7.57820969e+01 +84.29112j
     8.72257385e+01-154.79922j  ... -1.00665627e+02 -92.76081j
    -3.25569763e+01-143.90318j  -2.27663406e+02 +70.8449j  ]
   [ 2.59337730e+01 +47.23597j  -6.46899033e+01 +98.40452j
     9.90038757e+01-118.64778j  ... -2.73922150e+02+190.94734j
     1.20167503e+02-118.66812j   1.37966995e+02 -27.436005j]]]], shape=(1, 1, 686, 686), dtype=complex64)
tf.Tensor(
[[[[ 2.35175906e+05  +0.j       -1.30122070e+02 +71.79061j
     2.00236359e+02 -40.830986j ...  2.86872749e+01-174.89594j
     2.00236313e+02 +40.831j    -1.30122101e+02 -71.790565j]
   [ 2.59338074e+01 -47.23604j   1.37966949e+02 +27.436068j
     1.20167496e+02+118.6681j   ...  1.09129379e+02 -69.504425j
     9.90038986e+01+118.647736j -6.46898956e+01 -98.40455j ]
   [ 1.74128876e+02+181.25378j  -2.27663452e+02 -70.84499j
    -3.25569305e+01+143.9032j   ...  1.96747189e+01-331.84482j
     8.72258377e+01+154.79916j   7.57820816e+01 -84.29112j ]
   ...
   [-8.93064423e+01 -94.69147j   2.13672241e+02+246.61185j
     2.80917282e+01-143.79309j  ...  2.39895905e+02 +28.519115j
     9.76210938e+01 -32.177147j  5.52456627e+01-292.4319j  ]
   [ 1.74128815e+02-181.2538j    7.57820969e+01 +84.29112j
     8.72257385e+01-154.79922j  ... -1.00665627e+02 -92.76081j
    -3.25569763e+01-143.90318j  -2.27663406e+02 +70.8449j  ]
   [ 2.59337730e+01 +47.23597j  -6.46899033e+01 +98.40452j
     9.90038757e+01-118.64778j  ... -2.73922150e+02+190.94734j
     1.20167503e+02-118.66812j   1.37966995e+02 -27.436005j]]]], shape=(1, 1, 686, 686), dtype=complex64)
```

Output with TF2.12:
```
tf.Tensor(
[[[[ 2.3510438e+05  +0.j       -2.3397424e+02 +95.00578j
    -2.8460379e+02 +96.456604j ...  3.8885651e+00+184.90929j
    -2.8460379e+02 -96.45657j  -2.3397423e+02 -95.00578j ]
   [ 4.0683159e+01-196.26071j  -3.3500740e+01+253.18457j
     1.2614692e+02+251.45973j  ...  1.1127013e+02 +73.44855j
    -9.9558403e+01+301.1387j   -3.7669452e+02+258.1033j  ]
   [-2.9365808e+02 -26.070122j  1.1248355e+02+219.34619j
    -4.5030090e+01-103.50034j  ...  6.4451477e+01-142.80096j
    -2.9005927e+02 -39.51029j  -1.2654387e+02 +66.92355j ]
   ...
   [-5.9005920e+01-109.54649j  -2.5984564e+02  -8.701363j
    -1.3211296e+02 -64.48026j  ...  1.5030142e+02+232.03638j
    -2.1768060e+02 -23.536972j -4.0744392e+01-186.61447j ]
   [-2.9365802e+02 +26.070091j -1.2654389e+02 -66.92351j
    -2.9005920e+02 +39.51022j  ...  1.3489523e+02 -29.27546j
    -4.5030064e+01+103.50032j   1.1248353e+02-219.34612j ]
   [ 4.0683189e+01+196.26067j  -3.7669446e+02-258.1033j
    -9.9558449e+01-301.1387j   ...  9.9753399e+00 +51.120373j
     1.2614685e+02-251.45963j  -3.3500755e+01-253.18451j ]]]], shape=(1, 1, 686, 686), dtype=complex64)
tf.Tensor(
[[[[0.+0.j 0.+0.j 0.+0.j ... 0.+0.j 0.+0.j 0.+0.j]
   [0.+0.j 0.+0.j 0.+0.j ... 0.+0.j 0.+0.j 0.+0.j]
   [0.+0.j 0.+0.j 0.+0.j ... 0.+0.j 0.+0.j 0.+0.j]
   ...
   [0.+0.j 0.+0.j 0.+0.j ... 0.+0.j 0.+0.j 0.+0.j]
   [0.+0.j 0.+0.j 0.+0.j ... 0.+0.j 0.+0.j 0.+0.j]
   [0.+0.j 0.+0.j 0.+0.j ... 0.+0.j 0.+0.j 0.+0.j]]]], shape=(1, 1, 686, 686), dtype=complex64)

``` @sachinprasadhs any update here?  @reedwm is there any update for this issue? Thanks! @hubingallin started looking into this issue. A bisect shows 877e6ecbd43957d8e7313efb61e962149f6231cb (cl/505171308) as the culrpit. I'm guessing the issue is that the FFT plan cache does not include the device in the cache key (probably cufftHandle is associated with the current context's device, or maybe even a stream). So running the same FFT on multiple devices would reuse the cufftHandle

@cantonios can you fix this? Looks like @reedwm was right, it's because the `cufftHandle` is created on one device, cached, then attempted to be used on another physical device.  We can add the device ID to the plan cache key to address this.  The fix should be out soon. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60926">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60926">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

binary

### Tensorflow Version

2.12.0 2.14.0-dev20230619

### Custom Code

No

### OS Platform and Distribution

Linux Ubuntu 18.04

### Mobile device

_No response_

### Python version

3.9.16

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

11.8.0 / 8.6.0.163

### GPU model and memory

_No response_

### Current Behaviour?

Using TensorFlow FFT in a Keras model will produce incorrect results when using MirroredStrategy and multiple GPUs.
This is not an accuracy issue. The results of consecutive calls seem to be either correct or garbage.

I created a test Keras model that has one layer that does FFT. There is also a reference model using a DFT layer that is used to verify that incorrect behavior only happens when using tf.signal.fft.
Attached is a test application that runs both models in different combinations of MirroredStrategy/default strategy and eager/graph execution.
MirroredStrategy and graph execution is the combination that produces the error. At least two GPUs are required to reproduce the problem.
The output MAE loss is around 6.5, which translates to 650% error. (The absolute value of each entry in the correct output is 1.0.)

I think it's not a user error, but if it is, there should be an error or warning instead of incorrect results.

I was able to reproduce the issue with all TF fft variants (tf.signal.fft, tf.signal.rfft, tf.signal.stft, tf.signal.fft2d)

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import numpy as np
from scipy.linalg import dft
from math import sqrt

# Layer that does tf.signal.fft operation
class FFTLayer(tf.keras.layers.Layer):

    def call(self, x):
      fx = tf.signal.fft(x)
      return fx

# Layer that returns same results as tf.signal.fft op, but
# uses slower direct computation of DFT, implemented as matrix multiply.
class MatrixDFTLayer(tf.keras.layers.Layer):
    def __init__(self):
      super().__init__()
      self.dft = tf.cast(dft(1024), tf.complex64)

    def call(self, x):
        fx = self.dft @ tf.transpose(x)
        return tf.transpose(fx)


def create_model(use_mirrored_strategy: bool = True,
                              run_eagerly: bool = True,
                              layer_to_use: tf.keras.layers.Layer = FFTLayer) -> None:
    print(f"\ncreate model with: use_mirrored_strategy: {use_mirrored_strategy}, ",
          f"run_eagerly: {run_eagerly}, ",
          f"layer_to_use: {layer_to_use}")

    if use_mirrored_strategy:
        distribution_strategy = tf.distribute.MirroredStrategy()
    else:
        distribution_strategy = tf.distribute.get_strategy()
    with distribution_strategy.scope():

        ins = tf.keras.layers.Input([1024], dtype=tf.complex64)
        x = layer_to_use()(ins)
        model = tf.keras.Model(inputs=ins, outputs=x)

        model.compile(
            loss=tf.keras.losses.MeanAbsoluteError(),
            run_eagerly=run_eagerly
        )
    return model


def create_data(fft_size, batch_size, num_steps):
    num_examples = num_steps * batch_size

    # y data is a complex vector of all (1/sqrt(2), (1/sqrt(2)j)
    train_y = np.ones([fft_size], np.float32)
    train_y = (1/sqrt(2))*train_y + (1/sqrt(2))*1j*train_y
    # abs mean is 1 -> MAE magnitude should be compared to 1
    print("train_y mean: ", tf.reduce_mean(tf.abs(train_y)))

    # use inverse transform to create input data
    # fft(train_x) will produce train_y
    train_x = tf.signal.ifft(train_y)

    # clone data to get larger training set
    train_y = train_y[tf.newaxis, ...]
    train_x = train_x[tf.newaxis, ...]

    train_x = tf.tile(train_x, [num_examples, 1])
    train_y = tf.tile(train_y, [num_examples, 1])
    
    return train_x, train_y


fft_size = 1024
batch_size = 9
num_steps = 100
train_x, train_y = create_data(fft_size, batch_size, num_steps)

# Test cases with MatrixDFTLayer
# These are all ok, MAE close to 0.0
# ok
model = create_model(use_mirrored_strategy=False, run_eagerly=False, layer_to_use=MatrixDFTLayer)
loss = model.evaluate(train_x, train_y, batch_size=batch_size, verbose=0)
print(f"loss: {loss}")
# ok
model = create_model(use_mirrored_strategy=False, run_eagerly=True, layer_to_use=MatrixDFTLayer)
loss = model.evaluate(train_x, train_y, batch_size=batch_size, verbose=0)
print(f"loss: {loss}")
# ok
model = create_model(use_mirrored_strategy=True, run_eagerly=False, layer_to_use=MatrixDFTLayer)
loss = model.evaluate(train_x, train_y, batch_size=batch_size, verbose=0)
print(f"loss: {loss}")

# Test Cases using TF FFT. These fail when using MirroredStrategy.
# ok
model = create_model(use_mirrored_strategy=False, run_eagerly=False, layer_to_use=FFTLayer)
loss = model.evaluate(train_x, train_y, batch_size=batch_size, verbose=0)
print(f"loss: {loss}")
# ok
model = create_model(use_mirrored_strategy=False, run_eagerly=True, layer_to_use=FFTLayer)
loss = model.evaluate(train_x, train_y, batch_size=batch_size, verbose=0)
print(f"loss: {loss}")
# fail, 
model = create_model(use_mirrored_strategy=True, run_eagerly=False,layer_to_use= FFTLayer)
loss = model.evaluate(train_x, train_y, batch_size=batch_size, verbose=0)
print(f"loss: {loss}")
```


### Relevant log output

```shell
train_y mean:  tf.Tensor(1.0, shape=(), dtype=float32)

create model with: use_mirrored_strategy: False,  run_eagerly: False,  layer_to_use: <class '__main__.MatrixDFTLayer'>
loss: 0.0

create model with: use_mirrored_strategy: False,  run_eagerly: True,  layer_to_use: <class '__main__.MatrixDFTLayer'>
loss: 0.0

create model with: use_mirrored_strategy: True,  run_eagerly: False,  layer_to_use: <class '__main__.MatrixDFTLayer'>
loss: 0.0

create model with: use_mirrored_strategy: False,  run_eagerly: False,  layer_to_use: <class '__main__.FFTLayer'>
loss: 0.0

create model with: use_mirrored_strategy: False,  run_eagerly: True,  layer_to_use: <class '__main__.FFTLayer'>
loss: 0.0

create model with: use_mirrored_strategy: True,  run_eagerly: False,  layer_to_use: <class '__main__.FFTLayer'>
loss: 6.451958656311035
```
</details>