Custom shuffle layer leaks memory when run on Apple M1 GPU with `tensorflow-metal`
Hi @sirno ,

First thing I want to clear that `tensorflow-macos` was built and maintained by Apple itself . 

Hence I tried to test the code first with Regular Tensorflow package `tensorflow`, for confirming the memory leakage problem exists with TF package also.I have executed the code on colab with GPU environment and observed no memory leakage(ran upto 175 epochs) and attached [gist](https://colab.research.google.com/gist/SuryanarayanaY/767e352ddc144f0fa716e68eef7d9f6c/60616.ipynb) here for reference.

This indicates the issue is specific to only `tensorflow-macos` package and shall be addressed by Apple tensorflow-metal developes. You can post the issue [here](https://developer.apple.com/forums/tags/tensorflow-metal) .

Thanks! thanks for the response and the reference to the Apple forums.

I will resubmit the issue in the Apple forums. 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/60616">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/60616">No</a>

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

Bug

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

No

### Source

binary

### Tensorflow Version

2.12.0

### Custom Code

Yes

### OS Platform and Distribution

macOS 13.0 (22A380)

### Mobile device

Apple M1

### Python version

3.10.10

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

The following layer leaks memory during training with keras when running on an Apple M1 with `tensorflow-macos` and `tensorflow-metal` installed:

```
class Shuffle(keras.layers.Layer):
    def call(self, inputs):
        shape = tf.concat([tf.shape(inputs)[:-1], [1]], axis=0)
        rnd = tf.argsort(tf.random.uniform(shape), axis=1)
        return tf.gather_nd(inputs, rnd, batch_dims=2)
```

When run on the CPU alone it does not leak memory


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_datasets as tfds
from tensorflow_datasets.testing.mocking import mock_data

import numpy as np


class Shuffle(keras.layers.Layer):
    def call(self, inputs):
        shape = tf.concat([tf.shape(inputs)[:-1], [1]], axis=0)
        rnd = tf.argsort(tf.random.uniform(shape), axis=1)
        return tf.gather_nd(inputs, rnd, batch_dims=2)


def build_leaky_model(input_shape):
    input = keras.Input(input_shape)

    x = Shuffle()(input)

    x = layers.GlobalAveragePooling2D()(x)

    x = layers.Flatten()(x)

    output = layers.Dense(1, activation="sigmoid")(x)

    return keras.Model(
        input,
        output,
    )


epochs = 1000

data_set = tf.data.Dataset.from_generator(
    lambda: (
        (
            np.ones(shape=(1000, 5000, 1), dtype=np.uint8),
            i % 2,
        )
        for i in range(200)
    ),
    output_types=(tf.uint8, tf.int32),
    output_shapes=((1000, 5000, 1), ()),
)

data_set = data_set.batch(10)

model = build_leaky_model((1000, 5000, 1))

model.compile(
    optimizer=keras.optimizers.Adam(1e-3),
    loss="binary_crossentropy",
    metrics=["accuracy"],
)

history = model.fit(
    data_set,
    epochs=epochs,
)
```
```


### Relevant log output

_No response_</details>