Repository: tensorflow/addons

Issue:
Memory leak in cutout function
**Describe the bug**

There is a memory leak in the cutout function when using random sizes and offset and XLA.

**Code to reproduce the issue**

Run that on Collab with a GPU.
```python
import tensorflow as tf
import tensorflow_addons as tfa
import psutil

tf.config.optimizer.set_jit(True)
for gpu in tf.config.list_physical_devices("GPU"):
    tf.config.experimental.set_memory_growth(gpu, True)
process = psutil.Process()

image_size = 3000
@tf.function()
def random_cutout(img):
    img = tf.tile(img[tf.newaxis, ...], [16, 1, 1, 1])
    mask_size = tf.random.uniform(
        shape=[], minval=image_size // 200, maxval=image_size // 4, dtype=tf.int32
    )
    mask_size *= 2
    offset = tf.random.uniform(
        shape=[tf.shape(img)[0], 2],
        minval=0,
        maxval=image_size - mask_size,
        dtype=tf.int32,
    )
    offset += tf.expand_dims(mask_size, axis=0) // 2
    return tfa.image.cutout(img, tf.stack([mask_size, mask_size]), offset=offset)

img = tf.ones([image_size, image_size, 3])
for _ in range(100):
    for _ in range(1000):
        random_cutout(img)
    print(f"{process.memory_info().rss / 1024 ** 3:.2f} GB")
```

Outputs :
```
WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
1.48 GB
1.52 GB
1.54 GB
1.62 GB
1.64 GB
1.67 GB
1.69 GB
1.82 GB
1.84 GB
1.86 GB
1.89 GB
1.91 GB
1.94 GB
1.97 GB
2.18 GB
2.21 GB
```

**Cause**

This in function definition with outside tensors (lower_pads, ...) is causing the issue. Moving this function to the [cutout_ops.py ](https://github.com/tensorflow/addons/blob/master/tensorflow_addons/image/cutout_ops.py) namespace with `cutout_shape, padding_dims` as args fix the issue

[Code](https://github.com/tensorflow/addons/blob/104e345aaa0ef146e5e88fb5000e957db6856abf/tensorflow_addons/image/cutout_ops.py#L154) line 154 of cutout_ops.py.
```
        def fn(i):
            padding_dims = [
                [lower_pads[i], upper_pads[i]],
                [left_pads[i], right_pads[i]],
            ]
            mask = tf.pad(
                tf.zeros(cutout_shape[i], dtype=tf.bool),
                padding_dims,
                constant_values=True,
            )
            return mask

        mask = tf.map_fn(
            fn,
            tf.range(tf.shape(cutout_shape)[0]),
            fn_output_signature=tf.TensorSpec(
                shape=image_static_shape[1:-1], dtype=tf.bool
            ),
        )     
```

If needed (I don't know since tensorflow-addons is winding down), I can make a small merge request to fix this bug.


Discussion:
As cutout is already in keras_cv I suggest a migration. But if it is a small PR please submit it. https://keras.io/api/keras_cv/layers/preprocessing/random_cutout/

Resolving patch:
diff --git a/tensorflow_addons/image/cutout_ops.py b/tensorflow_addons/image/cutout_ops.py
--- a/tensorflow_addons/image/cutout_ops.py
+++ b/tensorflow_addons/image/cutout_ops.py
@@ -150,22 +150,16 @@ def cutout(
             ],
             [1, 0],
         )
-
-        def fn(i):
-            padding_dims = [
-                [lower_pads[i], upper_pads[i]],
-                [left_pads[i], right_pads[i]],
-            ]
-            mask = tf.pad(
-                tf.zeros(cutout_shape[i], dtype=tf.bool),
-                padding_dims,
-                constant_values=True,
-            )
-            return mask
-
+        padding_dims = tf.stack(
+            [
+                tf.stack([lower_pads, upper_pads], axis=1),
+                tf.stack([left_pads, right_pads], axis=1),
+            ],
+            axis=1,
+        )
         mask = tf.map_fn(
-            fn,
-            tf.range(tf.shape(cutout_shape)[0]),
+            _generate_masks,
+            [cutout_shape, padding_dims],
             fn_output_signature=tf.TensorSpec(
                 shape=image_static_shape[1:-1], dtype=tf.bool
             ),
@@ -180,3 +174,13 @@ def fn(i):
         )
         images.set_shape(image_static_shape)
         return images
+
+
+def _generate_masks(args):
+    cutout_shape, padding_dims = args
+    mask = tf.pad(
+        tf.zeros(cutout_shape, dtype=tf.bool),
+        padding_dims,
+        constant_values=True,
+    )
+    return mask


Test patch: