We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tensorflow_addons/image/cutout_ops.py`
Content:
```
1 # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 """Cutout op"""
16 
17 import tensorflow as tf
18 from tensorflow_addons.utils.types import TensorLike, Number
19 
20 
21 @tf.function
22 def _norm_params(mask_size, offset=None):
23     tf.assert_equal(
24         tf.reduce_any(mask_size % 2 != 0),
25         False,
26         "mask_size should be divisible by 2",
27     )
28     if tf.rank(mask_size) == 0:
29         mask_size = tf.stack([mask_size, mask_size])
30     if offset is not None and tf.rank(offset) == 1:
31         offset = tf.expand_dims(offset, 0)
32     return mask_size, offset
33 
34 
35 @tf.function
36 def _random_center(mask_dim_length, image_dim_length, batch_size, seed):
37     if mask_dim_length >= image_dim_length:
38         return tf.tile([image_dim_length // 2], [batch_size])
39     half_mask_dim_length = mask_dim_length // 2
40     return tf.random.uniform(
41         shape=[batch_size],
42         minval=half_mask_dim_length,
43         maxval=image_dim_length - half_mask_dim_length,
44         dtype=tf.int32,
45         seed=seed,
46     )
47 
48 
49 def random_cutout(
50     images: TensorLike,
51     mask_size: TensorLike,
52     constant_values: Number = 0,
53     seed: Number = None,
54 ) -> tf.Tensor:
55     """Apply [cutout](https://arxiv.org/abs/1708.04552) to images with random offset.
56 
57     This operation applies a `(mask_height x mask_width)` mask of zeros to
58     a random location within `images`. The pixel values filled in will be of
59     the value `constant_values`. The location where the mask will be applied is
60     randomly chosen uniformly over the whole images.
61 
62     Args:
63       images: A tensor of shape `(batch_size, height, width, channels)` (NHWC).
64       mask_size: Specifies how big the zero mask that will be generated is that
65         is applied to the images. The mask will be of size
66         `(mask_height x mask_width)`. Note: mask_size should be divisible by 2.
67       constant_values: What pixel value to fill in the images in the area that has
68         the cutout mask applied to it.
69       seed: A Python integer. Used in combination with `tf.random.set_seed` to
70         create a reproducible sequence of tensors across multiple calls.
71     Returns:
72       A `Tensor` of the same shape and dtype as `images`.
73     Raises:
74       InvalidArgumentError: if `mask_size` can't be divisible by 2.
75     """
76     images = tf.convert_to_tensor(images)
77     mask_size = tf.convert_to_tensor(mask_size)
78 
79     image_dynamic_shape = tf.shape(images)
80     batch_size, image_height, image_width = (
81         image_dynamic_shape[0],
82         image_dynamic_shape[1],
83         image_dynamic_shape[2],
84     )
85 
86     mask_size, _ = _norm_params(mask_size, offset=None)
87 
88     cutout_center_height = _random_center(mask_size[0], image_height, batch_size, seed)
89     cutout_center_width = _random_center(mask_size[1], image_width, batch_size, seed)
90 
91     offset = tf.transpose([cutout_center_height, cutout_center_width], [1, 0])
92     return cutout(images, mask_size, offset, constant_values)
93 
94 
95 def cutout(
96     images: TensorLike,
97     mask_size: TensorLike,
98     offset: TensorLike = (0, 0),
99     constant_values: Number = 0,
100 ) -> tf.Tensor:
101     """Apply [cutout](https://arxiv.org/abs/1708.04552) to images.
102 
103     This operation applies a `(mask_height x mask_width)` mask of zeros to
104     a location within `images` specified by the offset.
105     The pixel values filled in will be of the value `constant_values`.
106     The location where the mask will be applied is randomly
107     chosen uniformly over the whole images.
108 
109     Args:
110       images: A tensor of shape `(batch_size, height, width, channels)` (NHWC).
111       mask_size: Specifies how big the zero mask that will be generated is that
112         is applied to the images. The mask will be of size
113         `(mask_height x mask_width)`. Note: mask_size should be divisible by 2.
114       offset: A tuple of `(height, width)` or `(batch_size, 2)`
115       constant_values: What pixel value to fill in the images in the area that has
116         the cutout mask applied to it.
117     Returns:
118       A `Tensor` of the same shape and dtype as `images`.
119     Raises:
120       InvalidArgumentError: if `mask_size` can't be divisible by 2.
121     """
122     with tf.name_scope("cutout"):
123         images = tf.convert_to_tensor(images)
124         mask_size = tf.convert_to_tensor(mask_size)
125         offset = tf.convert_to_tensor(offset)
126 
127         image_static_shape = images.shape
128         image_dynamic_shape = tf.shape(images)
129         image_height, image_width, channels = (
130             image_dynamic_shape[1],
131             image_dynamic_shape[2],
132             image_dynamic_shape[3],
133         )
134 
135         mask_size, offset = _norm_params(mask_size, offset)
136         mask_size = mask_size // 2
137 
138         cutout_center_heights = offset[:, 0]
139         cutout_center_widths = offset[:, 1]
140 
141         lower_pads = tf.maximum(0, cutout_center_heights - mask_size[0])
142         upper_pads = tf.maximum(0, image_height - cutout_center_heights - mask_size[0])
143         left_pads = tf.maximum(0, cutout_center_widths - mask_size[1])
144         right_pads = tf.maximum(0, image_width - cutout_center_widths - mask_size[1])
145 
146         cutout_shape = tf.transpose(
147             [
148                 image_height - (lower_pads + upper_pads),
149                 image_width - (left_pads + right_pads),
150             ],
151             [1, 0],
152         )
153 
154         def fn(i):
155             padding_dims = [
156                 [lower_pads[i], upper_pads[i]],
157                 [left_pads[i], right_pads[i]],
158             ]
159             mask = tf.pad(
160                 tf.zeros(cutout_shape[i], dtype=tf.bool),
161                 padding_dims,
162                 constant_values=True,
163             )
164             return mask
165 
166         mask = tf.map_fn(
167             fn,
168             tf.range(tf.shape(cutout_shape)[0]),
169             fn_output_signature=tf.TensorSpec(
170                 shape=image_static_shape[1:-1], dtype=tf.bool
171             ),
172         )
173         mask = tf.expand_dims(mask, -1)
174         mask = tf.tile(mask, [1, 1, 1, channels])
175 
176         images = tf.where(
177             mask,
178             images,
179             tf.cast(constant_values, dtype=images.dtype),
180         )
181         images.set_shape(image_static_shape)
182         return images
183
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
             ],
             [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 @@
         )
         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

{"golden_diff": "diff --git a/tensorflow_addons/image/cutout_ops.py b/tensorflow_addons/image/cutout_ops.py\n--- a/tensorflow_addons/image/cutout_ops.py\n+++ b/tensorflow_addons/image/cutout_ops.py\n@@ -150,22 +150,16 @@\n             ],\n             [1, 0],\n         )\n-\n-        def fn(i):\n-            padding_dims = [\n-                [lower_pads[i], upper_pads[i]],\n-                [left_pads[i], right_pads[i]],\n-            ]\n-            mask = tf.pad(\n-                tf.zeros(cutout_shape[i], dtype=tf.bool),\n-                padding_dims,\n-                constant_values=True,\n-            )\n-            return mask\n-\n+        padding_dims = tf.stack(\n+            [\n+                tf.stack([lower_pads, upper_pads], axis=1),\n+                tf.stack([left_pads, right_pads], axis=1),\n+            ],\n+            axis=1,\n+        )\n         mask = tf.map_fn(\n-            fn,\n-            tf.range(tf.shape(cutout_shape)[0]),\n+            _generate_masks,\n+            [cutout_shape, padding_dims],\n             fn_output_signature=tf.TensorSpec(\n                 shape=image_static_shape[1:-1], dtype=tf.bool\n             ),\n@@ -180,3 +174,13 @@\n         )\n         images.set_shape(image_static_shape)\n         return images\n+\n+\n+def _generate_masks(args):\n+    cutout_shape, padding_dims = args\n+    mask = tf.pad(\n+        tf.zeros(cutout_shape, dtype=tf.bool),\n+        padding_dims,\n+        constant_values=True,\n+    )\n+    return mask\n", "issue": "Memory leak in cutout function\n**Describe the bug**\r\n\r\nThere is a memory leak in the cutout function when using random sizes and offset and XLA.\r\n\r\n**Code to reproduce the issue**\r\n\r\nRun that on Collab with a GPU.\r\n```python\r\nimport tensorflow as tf\r\nimport tensorflow_addons as tfa\r\nimport psutil\r\n\r\ntf.config.optimizer.set_jit(True)\r\nfor gpu in tf.config.list_physical_devices(\"GPU\"):\r\n    tf.config.experimental.set_memory_growth(gpu, True)\r\nprocess = psutil.Process()\r\n\r\nimage_size = 3000\r\n@tf.function()\r\ndef random_cutout(img):\r\n    img = tf.tile(img[tf.newaxis, ...], [16, 1, 1, 1])\r\n    mask_size = tf.random.uniform(\r\n        shape=[], minval=image_size // 200, maxval=image_size // 4, dtype=tf.int32\r\n    )\r\n    mask_size *= 2\r\n    offset = tf.random.uniform(\r\n        shape=[tf.shape(img)[0], 2],\r\n        minval=0,\r\n        maxval=image_size - mask_size,\r\n        dtype=tf.int32,\r\n    )\r\n    offset += tf.expand_dims(mask_size, axis=0) // 2\r\n    return tfa.image.cutout(img, tf.stack([mask_size, mask_size]), offset=offset)\r\n\r\nimg = tf.ones([image_size, image_size, 3])\r\nfor _ in range(100):\r\n    for _ in range(1000):\r\n        random_cutout(img)\r\n    print(f\"{process.memory_info().rss / 1024 ** 3:.2f} GB\")\r\n```\r\n\r\nOutputs :\r\n```\r\nWARNING: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.\r\nInstructions for updating:\r\nLambda 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\r\n1.48 GB\r\n1.52 GB\r\n1.54 GB\r\n1.62 GB\r\n1.64 GB\r\n1.67 GB\r\n1.69 GB\r\n1.82 GB\r\n1.84 GB\r\n1.86 GB\r\n1.89 GB\r\n1.91 GB\r\n1.94 GB\r\n1.97 GB\r\n2.18 GB\r\n2.21 GB\r\n```\r\n\r\n**Cause**\r\n\r\nThis 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\r\n\r\n[Code](https://github.com/tensorflow/addons/blob/104e345aaa0ef146e5e88fb5000e957db6856abf/tensorflow_addons/image/cutout_ops.py#L154) line 154 of cutout_ops.py.\r\n```\r\n        def fn(i):\r\n            padding_dims = [\r\n                [lower_pads[i], upper_pads[i]],\r\n                [left_pads[i], right_pads[i]],\r\n            ]\r\n            mask = tf.pad(\r\n                tf.zeros(cutout_shape[i], dtype=tf.bool),\r\n                padding_dims,\r\n                constant_values=True,\r\n            )\r\n            return mask\r\n\r\n        mask = tf.map_fn(\r\n            fn,\r\n            tf.range(tf.shape(cutout_shape)[0]),\r\n            fn_output_signature=tf.TensorSpec(\r\n                shape=image_static_shape[1:-1], dtype=tf.bool\r\n            ),\r\n        )     \r\n```\r\n\r\nIf needed (I don't know since tensorflow-addons is winding down), I can make a small merge request to fix this bug.\n", "before_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Cutout op\"\"\"\n\nimport tensorflow as tf\nfrom tensorflow_addons.utils.types import TensorLike, Number\n\n\n@tf.function\ndef _norm_params(mask_size, offset=None):\n    tf.assert_equal(\n        tf.reduce_any(mask_size % 2 != 0),\n        False,\n        \"mask_size should be divisible by 2\",\n    )\n    if tf.rank(mask_size) == 0:\n        mask_size = tf.stack([mask_size, mask_size])\n    if offset is not None and tf.rank(offset) == 1:\n        offset = tf.expand_dims(offset, 0)\n    return mask_size, offset\n\n\n@tf.function\ndef _random_center(mask_dim_length, image_dim_length, batch_size, seed):\n    if mask_dim_length >= image_dim_length:\n        return tf.tile([image_dim_length // 2], [batch_size])\n    half_mask_dim_length = mask_dim_length // 2\n    return tf.random.uniform(\n        shape=[batch_size],\n        minval=half_mask_dim_length,\n        maxval=image_dim_length - half_mask_dim_length,\n        dtype=tf.int32,\n        seed=seed,\n    )\n\n\ndef random_cutout(\n    images: TensorLike,\n    mask_size: TensorLike,\n    constant_values: Number = 0,\n    seed: Number = None,\n) -> tf.Tensor:\n    \"\"\"Apply [cutout](https://arxiv.org/abs/1708.04552) to images with random offset.\n\n    This operation applies a `(mask_height x mask_width)` mask of zeros to\n    a random location within `images`. The pixel values filled in will be of\n    the value `constant_values`. The location where the mask will be applied is\n    randomly chosen uniformly over the whole images.\n\n    Args:\n      images: A tensor of shape `(batch_size, height, width, channels)` (NHWC).\n      mask_size: Specifies how big the zero mask that will be generated is that\n        is applied to the images. The mask will be of size\n        `(mask_height x mask_width)`. Note: mask_size should be divisible by 2.\n      constant_values: What pixel value to fill in the images in the area that has\n        the cutout mask applied to it.\n      seed: A Python integer. Used in combination with `tf.random.set_seed` to\n        create a reproducible sequence of tensors across multiple calls.\n    Returns:\n      A `Tensor` of the same shape and dtype as `images`.\n    Raises:\n      InvalidArgumentError: if `mask_size` can't be divisible by 2.\n    \"\"\"\n    images = tf.convert_to_tensor(images)\n    mask_size = tf.convert_to_tensor(mask_size)\n\n    image_dynamic_shape = tf.shape(images)\n    batch_size, image_height, image_width = (\n        image_dynamic_shape[0],\n        image_dynamic_shape[1],\n        image_dynamic_shape[2],\n    )\n\n    mask_size, _ = _norm_params(mask_size, offset=None)\n\n    cutout_center_height = _random_center(mask_size[0], image_height, batch_size, seed)\n    cutout_center_width = _random_center(mask_size[1], image_width, batch_size, seed)\n\n    offset = tf.transpose([cutout_center_height, cutout_center_width], [1, 0])\n    return cutout(images, mask_size, offset, constant_values)\n\n\ndef cutout(\n    images: TensorLike,\n    mask_size: TensorLike,\n    offset: TensorLike = (0, 0),\n    constant_values: Number = 0,\n) -> tf.Tensor:\n    \"\"\"Apply [cutout](https://arxiv.org/abs/1708.04552) to images.\n\n    This operation applies a `(mask_height x mask_width)` mask of zeros to\n    a location within `images` specified by the offset.\n    The pixel values filled in will be of the value `constant_values`.\n    The location where the mask will be applied is randomly\n    chosen uniformly over the whole images.\n\n    Args:\n      images: A tensor of shape `(batch_size, height, width, channels)` (NHWC).\n      mask_size: Specifies how big the zero mask that will be generated is that\n        is applied to the images. The mask will be of size\n        `(mask_height x mask_width)`. Note: mask_size should be divisible by 2.\n      offset: A tuple of `(height, width)` or `(batch_size, 2)`\n      constant_values: What pixel value to fill in the images in the area that has\n        the cutout mask applied to it.\n    Returns:\n      A `Tensor` of the same shape and dtype as `images`.\n    Raises:\n      InvalidArgumentError: if `mask_size` can't be divisible by 2.\n    \"\"\"\n    with tf.name_scope(\"cutout\"):\n        images = tf.convert_to_tensor(images)\n        mask_size = tf.convert_to_tensor(mask_size)\n        offset = tf.convert_to_tensor(offset)\n\n        image_static_shape = images.shape\n        image_dynamic_shape = tf.shape(images)\n        image_height, image_width, channels = (\n            image_dynamic_shape[1],\n            image_dynamic_shape[2],\n            image_dynamic_shape[3],\n        )\n\n        mask_size, offset = _norm_params(mask_size, offset)\n        mask_size = mask_size // 2\n\n        cutout_center_heights = offset[:, 0]\n        cutout_center_widths = offset[:, 1]\n\n        lower_pads = tf.maximum(0, cutout_center_heights - mask_size[0])\n        upper_pads = tf.maximum(0, image_height - cutout_center_heights - mask_size[0])\n        left_pads = tf.maximum(0, cutout_center_widths - mask_size[1])\n        right_pads = tf.maximum(0, image_width - cutout_center_widths - mask_size[1])\n\n        cutout_shape = tf.transpose(\n            [\n                image_height - (lower_pads + upper_pads),\n                image_width - (left_pads + right_pads),\n            ],\n            [1, 0],\n        )\n\n        def fn(i):\n            padding_dims = [\n                [lower_pads[i], upper_pads[i]],\n                [left_pads[i], right_pads[i]],\n            ]\n            mask = tf.pad(\n                tf.zeros(cutout_shape[i], dtype=tf.bool),\n                padding_dims,\n                constant_values=True,\n            )\n            return mask\n\n        mask = tf.map_fn(\n            fn,\n            tf.range(tf.shape(cutout_shape)[0]),\n            fn_output_signature=tf.TensorSpec(\n                shape=image_static_shape[1:-1], dtype=tf.bool\n            ),\n        )\n        mask = tf.expand_dims(mask, -1)\n        mask = tf.tile(mask, [1, 1, 1, channels])\n\n        images = tf.where(\n            mask,\n            images,\n            tf.cast(constant_values, dtype=images.dtype),\n        )\n        images.set_shape(image_static_shape)\n        return images\n", "path": "tensorflow_addons/image/cutout_ops.py"}], "after_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Cutout op\"\"\"\n\nimport tensorflow as tf\nfrom tensorflow_addons.utils.types import TensorLike, Number\n\n\n@tf.function\ndef _norm_params(mask_size, offset=None):\n    tf.assert_equal(\n        tf.reduce_any(mask_size % 2 != 0),\n        False,\n        \"mask_size should be divisible by 2\",\n    )\n    if tf.rank(mask_size) == 0:\n        mask_size = tf.stack([mask_size, mask_size])\n    if offset is not None and tf.rank(offset) == 1:\n        offset = tf.expand_dims(offset, 0)\n    return mask_size, offset\n\n\n@tf.function\ndef _random_center(mask_dim_length, image_dim_length, batch_size, seed):\n    if mask_dim_length >= image_dim_length:\n        return tf.tile([image_dim_length // 2], [batch_size])\n    half_mask_dim_length = mask_dim_length // 2\n    return tf.random.uniform(\n        shape=[batch_size],\n        minval=half_mask_dim_length,\n        maxval=image_dim_length - half_mask_dim_length,\n        dtype=tf.int32,\n        seed=seed,\n    )\n\n\ndef random_cutout(\n    images: TensorLike,\n    mask_size: TensorLike,\n    constant_values: Number = 0,\n    seed: Number = None,\n) -> tf.Tensor:\n    \"\"\"Apply [cutout](https://arxiv.org/abs/1708.04552) to images with random offset.\n\n    This operation applies a `(mask_height x mask_width)` mask of zeros to\n    a random location within `images`. The pixel values filled in will be of\n    the value `constant_values`. The location where the mask will be applied is\n    randomly chosen uniformly over the whole images.\n\n    Args:\n      images: A tensor of shape `(batch_size, height, width, channels)` (NHWC).\n      mask_size: Specifies how big the zero mask that will be generated is that\n        is applied to the images. The mask will be of size\n        `(mask_height x mask_width)`. Note: mask_size should be divisible by 2.\n      constant_values: What pixel value to fill in the images in the area that has\n        the cutout mask applied to it.\n      seed: A Python integer. Used in combination with `tf.random.set_seed` to\n        create a reproducible sequence of tensors across multiple calls.\n    Returns:\n      A `Tensor` of the same shape and dtype as `images`.\n    Raises:\n      InvalidArgumentError: if `mask_size` can't be divisible by 2.\n    \"\"\"\n    images = tf.convert_to_tensor(images)\n    mask_size = tf.convert_to_tensor(mask_size)\n\n    image_dynamic_shape = tf.shape(images)\n    batch_size, image_height, image_width = (\n        image_dynamic_shape[0],\n        image_dynamic_shape[1],\n        image_dynamic_shape[2],\n    )\n\n    mask_size, _ = _norm_params(mask_size, offset=None)\n\n    cutout_center_height = _random_center(mask_size[0], image_height, batch_size, seed)\n    cutout_center_width = _random_center(mask_size[1], image_width, batch_size, seed)\n\n    offset = tf.transpose([cutout_center_height, cutout_center_width], [1, 0])\n    return cutout(images, mask_size, offset, constant_values)\n\n\ndef cutout(\n    images: TensorLike,\n    mask_size: TensorLike,\n    offset: TensorLike = (0, 0),\n    constant_values: Number = 0,\n) -> tf.Tensor:\n    \"\"\"Apply [cutout](https://arxiv.org/abs/1708.04552) to images.\n\n    This operation applies a `(mask_height x mask_width)` mask of zeros to\n    a location within `images` specified by the offset.\n    The pixel values filled in will be of the value `constant_values`.\n    The location where the mask will be applied is randomly\n    chosen uniformly over the whole images.\n\n    Args:\n      images: A tensor of shape `(batch_size, height, width, channels)` (NHWC).\n      mask_size: Specifies how big the zero mask that will be generated is that\n        is applied to the images. The mask will be of size\n        `(mask_height x mask_width)`. Note: mask_size should be divisible by 2.\n      offset: A tuple of `(height, width)` or `(batch_size, 2)`\n      constant_values: What pixel value to fill in the images in the area that has\n        the cutout mask applied to it.\n    Returns:\n      A `Tensor` of the same shape and dtype as `images`.\n    Raises:\n      InvalidArgumentError: if `mask_size` can't be divisible by 2.\n    \"\"\"\n    with tf.name_scope(\"cutout\"):\n        images = tf.convert_to_tensor(images)\n        mask_size = tf.convert_to_tensor(mask_size)\n        offset = tf.convert_to_tensor(offset)\n\n        image_static_shape = images.shape\n        image_dynamic_shape = tf.shape(images)\n        image_height, image_width, channels = (\n            image_dynamic_shape[1],\n            image_dynamic_shape[2],\n            image_dynamic_shape[3],\n        )\n\n        mask_size, offset = _norm_params(mask_size, offset)\n        mask_size = mask_size // 2\n\n        cutout_center_heights = offset[:, 0]\n        cutout_center_widths = offset[:, 1]\n\n        lower_pads = tf.maximum(0, cutout_center_heights - mask_size[0])\n        upper_pads = tf.maximum(0, image_height - cutout_center_heights - mask_size[0])\n        left_pads = tf.maximum(0, cutout_center_widths - mask_size[1])\n        right_pads = tf.maximum(0, image_width - cutout_center_widths - mask_size[1])\n\n        cutout_shape = tf.transpose(\n            [\n                image_height - (lower_pads + upper_pads),\n                image_width - (left_pads + right_pads),\n            ],\n            [1, 0],\n        )\n        padding_dims = tf.stack(\n            [\n                tf.stack([lower_pads, upper_pads], axis=1),\n                tf.stack([left_pads, right_pads], axis=1),\n            ],\n            axis=1,\n        )\n        mask = tf.map_fn(\n            _generate_masks,\n            [cutout_shape, padding_dims],\n            fn_output_signature=tf.TensorSpec(\n                shape=image_static_shape[1:-1], dtype=tf.bool\n            ),\n        )\n        mask = tf.expand_dims(mask, -1)\n        mask = tf.tile(mask, [1, 1, 1, channels])\n\n        images = tf.where(\n            mask,\n            images,\n            tf.cast(constant_values, dtype=images.dtype),\n        )\n        images.set_shape(image_static_shape)\n        return images\n\n\ndef _generate_masks(args):\n    cutout_shape, padding_dims = args\n    mask = tf.pad(\n        tf.zeros(cutout_shape, dtype=tf.bool),\n        padding_dims,\n        constant_values=True,\n    )\n    return mask\n", "path": "tensorflow_addons/image/cutout_ops.py"}]}