We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
torch.nn.functional.affine_grid


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `ivy/functional/frontends/torch/nn/functional/vision_functions.py`
Content:
```
1 # global
2 import math
3 
4 # local
5 import ivy
6 from ivy import with_unsupported_dtypes
7 from ivy.functional.frontends.torch.func_wrapper import to_ivy_arrays_and_back
8 from ivy.utils.exceptions import IvyNotImplementedException
9 
10 
11 @to_ivy_arrays_and_back
12 def pixel_shuffle(input, upscale_factor):
13     input_shape = ivy.shape(input)
14 
15     ivy.utils.assertions.check_equal(
16         ivy.get_num_dims(input),
17         4,
18         message="pixel_shuffle expects 4D input, but got input with sizes "
19         + str(input_shape),
20         as_array=False,
21     )
22     b = input_shape[0]
23     c = input_shape[1]
24     h = input_shape[2]
25     w = input_shape[3]
26     upscale_factor_squared = upscale_factor * upscale_factor
27     ivy.utils.assertions.check_equal(
28         c % upscale_factor_squared,
29         0,
30         message="pixel_shuffle expects input channel to be divisible by square "
31         + "of upscale_factor, but got input with sizes "
32         + str(input_shape)
33         + ", upscale_factor="
34         + str(upscale_factor)
35         + ", and self.size(1)="
36         + str(c)
37         + " is not divisible by "
38         + str(upscale_factor_squared),
39         as_array=False,
40     )
41     oc = int(c / upscale_factor_squared)
42     oh = h * upscale_factor
43     ow = w * upscale_factor
44 
45     input_reshaped = ivy.reshape(input, (b, oc, upscale_factor, upscale_factor, h, w))
46     return ivy.reshape(
47         ivy.permute_dims(input_reshaped, (0, 1, 4, 2, 5, 3)), (b, oc, oh, ow)
48     )
49 
50 
51 @to_ivy_arrays_and_back
52 def pixel_unshuffle(input, downscale_factor):
53     input_shape = ivy.shape(input)
54 
55     ivy.utils.assertions.check_equal(
56         ivy.get_num_dims(input),
57         4,
58         message=(
59             f"pixel_unshuffle expects 4D input, but got input with sizes {input_shape}"
60         ),
61         as_array=False,
62     ),
63 
64     b = input_shape[0]
65     c = input_shape[1]
66     h = input_shape[2]
67     w = input_shape[3]
68     downscale_factor_squared = downscale_factor * downscale_factor
69 
70     ivy.utils.assertions.check_equal(
71         [h % downscale_factor, w % downscale_factor],
72         [0, 0],  # Assert h % downscale_factor == 0 and w % downscale_factor == 0
73         message=(
74             "pixel_unshuffle expects input height and width to be divisible by "
75             f"downscale_factor, but got input with sizes {input_shape}"
76             f", downscale_factor= {downscale_factor}"
77             f", and either self.size(2)= {h}"
78             f" or self.size(3)= {w}"
79             f" is not divisible by {downscale_factor}"
80         ),
81         as_array=False,
82     )
83     oc = c * downscale_factor_squared
84     oh = int(h / downscale_factor)
85     ow = int(w / downscale_factor)
86 
87     input_reshaped = ivy.reshape(
88         input, (b, c, oh, downscale_factor, ow, downscale_factor)
89     )
90     return ivy.reshape(
91         ivy.permute_dims(input_reshaped, (0, 1, 3, 5, 2, 4)), (b, oc, oh, ow)
92     )
93 
94 
95 def _handle_padding_shape(padding, n, mode):
96     padding = tuple(
97         [
98             (padding[i * 2], padding[i * 2 + 1])
99             for i in range(int(len(padding) / 2) - 1, -1, -1)
100         ]
101     )
102     while len(padding) < n:
103         if mode == "circular":
104             padding = padding + ((0, 0),)
105         else:
106             padding = ((0, 0),) + padding
107     if mode == "circular":
108         padding = tuple(list(padding)[::-1])
109     return padding
110 
111 
112 @to_ivy_arrays_and_back
113 def pad(input, pad, mode="constant", value=0):
114     mode_dict = {
115         "constant": "constant",
116         "reflect": "reflect",
117         "replicate": "edge",
118         "circular": "wrap",
119     }
120     if mode not in mode_dict:
121         raise ValueError(f"Unsupported padding mode: {mode}")
122     pad = _handle_padding_shape(pad, len(input.shape), mode)
123     return ivy.pad(input, pad, mode=mode_dict[mode], constant_values=value)
124 
125 
126 def _get_new_width_height(w_old, h_old, size=None, scale_factor=None):
127     if scale_factor and (not size):
128         if type(scale_factor) == int:
129             h_new = int(w_old * scale_factor)
130             w_new = int(h_old * scale_factor)
131         elif type(scale_factor) == tuple:
132             h_new = int(w_old * scale_factor[0])
133             w_new = int(h_old * scale_factor[1])
134     elif (not scale_factor) and size:
135         if type(size) == int:
136             h_new = size
137             w_new = size
138         elif type(size) == tuple:
139             h_new, w_new = size
140     return h_new, w_new
141 
142 
143 @with_unsupported_dtypes(
144     {
145         "2.0.1 and below": (
146             "bfloat16",
147             "float16",
148         )
149     },
150     "torch",
151 )
152 @to_ivy_arrays_and_back
153 def interpolate(
154     input,
155     size=None,
156     scale_factor=None,
157     mode="nearest",
158     align_corners=None,
159     recompute_scale_factor=None,
160     antialias=False,
161 ):
162     if mode in ["nearest", "area", "nearest-exact"]:
163         ivy.utils.assertions.check_exists(
164             align_corners,
165             inverse=True,
166             message=(
167                 "align_corners option can only be set with the interpolating modes:"
168                 " linear | bilinear | bicubic | trilinear"
169             ),
170         )
171     else:
172         if not ivy.exists(align_corners):
173             align_corners = False
174 
175     dim = ivy.get_num_dims(input) - 2  # Number of spatial dimensions.
176 
177     if ivy.exists(size) and ivy.exists(scale_factor):
178         raise ivy.utils.exceptions.IvyException(
179             "only one of size or scale_factor should be defined"
180         )
181 
182     elif ivy.exists(size) and not ivy.exists(scale_factor):
183         scale_factors = None
184 
185         if isinstance(size, (list, tuple)):
186             ivy.utils.assertions.check_equal(
187                 len(size),
188                 dim,
189                 inverse=False,
190                 message=(
191                     "Input and output must have the "
192                     "same number of spatial dimensions,"
193                     f" but got input with spatial dimensions of {list(input.shape[2:])}"
194                     f" and output size of {size}. "
195                     "Please provide input tensor in (N, C, d1, d2, ...,dK) format"
196                     " and output size in (o1, o2, ...,oK) format."
197                 ),
198                 as_array=False,
199             )
200             output_size = size
201         else:
202             output_size = [size for _ in range(dim)]
203 
204     elif ivy.exists(scale_factor) and not ivy.exists(size):
205         output_size = None
206 
207         if isinstance(scale_factor, (list, tuple)):
208             ivy.utils.assertions.check_equal(
209                 len(scale_factor),
210                 dim,
211                 inverse=False,
212                 message=(
213                     "Input and scale_factor must have the "
214                     "same number of spatial dimensions,"
215                     f" but got input with spatial dimensions of {list(input.shape[2:])}"
216                     f" and scale_factor of shape {scale_factor}. "
217                     "Please provide input tensor in (N, C, d1, d2, ...,dK) format"
218                     " and scale_factor in (s1, s2, ...,sK) format."
219                 ),
220                 as_array=False,
221             )
222             scale_factors = scale_factor
223         else:
224             scale_factors = [scale_factor for _ in range(dim)]
225 
226     else:
227         ivy.utils.assertions.check_any(
228             [ivy.exists(size), ivy.exists(scale_factor)],
229             message="either size or scale_factor should be defined",
230             as_array=False,
231         )
232 
233     if (
234         ivy.exists(size)
235         and ivy.exists(recompute_scale_factor)
236         and bool(recompute_scale_factor)
237     ):
238         raise ivy.utils.exceptions.IvyException(
239             "recompute_scale_factor is not meaningful with an explicit size."
240         )
241 
242     if ivy.exists(scale_factors):
243         output_size = [
244             math.floor(ivy.shape(input)[i + 2] * scale_factors[i]) for i in range(dim)
245         ]
246 
247     if (
248         bool(antialias)
249         and not (mode in ["bilinear", "bicubic"])
250         and ivy.get_num_dims(input) == 4
251     ):
252         raise ivy.utils.exceptions.IvyException(
253             "recompute_scale_factor is not meaningful with an explicit size."
254         )
255 
256     if ivy.get_num_dims(input) == 3 and mode == "bilinear":
257         raise IvyNotImplementedException(
258             "Got 3D input, but bilinear mode needs 4D input"
259         )
260     if ivy.get_num_dims(input) == 3 and mode == "trilinear":
261         raise IvyNotImplementedException(
262             "Got 3D input, but trilinear mode needs 5D input"
263         )
264     if ivy.get_num_dims(input) == 4 and mode == "linear":
265         raise IvyNotImplementedException("Got 4D input, but linear mode needs 3D input")
266     if ivy.get_num_dims(input) == 4 and mode == "trilinear":
267         raise IvyNotImplementedException(
268             "Got 4D input, but trilinear mode needs 5D input"
269         )
270     if ivy.get_num_dims(input) == 5 and mode == "linear":
271         raise IvyNotImplementedException("Got 5D input, but linear mode needs 3D input")
272     if ivy.get_num_dims(input) == 5 and mode == "bilinear":
273         raise IvyNotImplementedException(
274             "Got 5D input, but bilinear mode needs 4D input"
275         )
276 
277     ivy.utils.assertions.check_elem_in_list(
278         ivy.get_num_dims(input),
279         range(3, 6),
280         message=(
281             "Input Error: Only 3D, 4D and 5D input Tensors supported (got"
282             f" {ivy.get_num_dims(input)}D) for the modes: nearest | linear | bilinear |"
283             f" bicubic | trilinear | area | nearest-exact (got {mode})"
284         ),
285     )
286 
287     return ivy.interpolate(
288         input, output_size, mode=mode, align_corners=align_corners, antialias=antialias
289     )
290 
291 
292 @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, "torch")
293 @to_ivy_arrays_and_back
294 def upsample(
295     input,
296     size=None,
297     scale_factor=None,
298     mode="nearest",
299     align_corners=None,
300 ):
301     return interpolate(
302         input,
303         size=size,
304         scale_factor=scale_factor,
305         mode=mode,
306         align_corners=align_corners,
307     )
308 
309 
310 @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, "torch")
311 @to_ivy_arrays_and_back
312 def upsample_nearest(input, size=None, scale_factor=None):
313     return interpolate(input, size=size, scale_factor=scale_factor, mode="nearest")
314 
315 
316 @with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, "torch")
317 @to_ivy_arrays_and_back
318 def upsample_bilinear(input, size=None, scale_factor=None):
319     return interpolate(
320         input, size=size, scale_factor=scale_factor, mode="bilinear", align_corners=True
321     )
322
```
--- 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/ivy/functional/frontends/torch/nn/functional/vision_functions.py b/ivy/functional/frontends/torch/nn/functional/vision_functions.py
--- a/ivy/functional/frontends/torch/nn/functional/vision_functions.py
+++ b/ivy/functional/frontends/torch/nn/functional/vision_functions.py
@@ -319,3 +319,48 @@
     return interpolate(
         input, size=size, scale_factor=scale_factor, mode="bilinear", align_corners=True
     )
+
+
+@with_unsupported_dtypes({"2.0.1 and below": ("float16", "bfloat16")}, "torch")
+@to_ivy_arrays_and_back
+def affine_grid(theta, size, align_corners=False):
+    if len(size) == 4:
+        N, C, H, W = size
+        base_grid = ivy.empty((N, H, W, 3))
+        if align_corners:
+            base_grid[:, :, :, 0] = ivy.linspace(-1, 1, W)
+            base_grid[:, :, :, 1] = ivy.expand_dims(ivy.linspace(-1, 1, H), axis=-1)
+            base_grid[:, :, :, 2] = ivy.full((H, W), 1)
+            grid = ivy.matmul(base_grid.view((N, H * W, 3)), theta.swapaxes(1, 2))
+            return grid.view((N, H, W, 2))
+        else:
+            base_grid[:, :, :, 0] = ivy.linspace(-1, 1, W) * (W - 1) / W
+            base_grid[:, :, :, 1] = ivy.expand_dims(
+                ivy.linspace(-1, 1, H) * (H - 1) / H, axis=-1
+            )
+            base_grid[:, :, :, 2] = ivy.full((H, W), 1)
+        grid = ivy.matmul(base_grid.view((N, H * W, 3)), ivy.swapaxes(theta, 1, 2))
+        return grid.view((N, H, W, 2))
+    else:
+        N, C, D, H, W = size
+        base_grid = ivy.empty((N, D, H, W, 4))
+        if align_corners:
+            base_grid[:, :, :, :, 0] = ivy.linspace(-1, 1, W)
+            base_grid[:, :, :, :, 1] = ivy.expand_dims(ivy.linspace(-1, 1, H), axis=-1)
+            base_grid[:, :, :, :, 2] = ivy.expand_dims(
+                ivy.expand_dims(ivy.linspace(-1, 1, D), axis=-1), axis=-1
+            )
+            base_grid[:, :, :, :, 3] = ivy.full((D, H, W), 1)
+            grid = ivy.matmul(base_grid.view((N, D * H * W, 4)), theta.swapaxes(1, 2))
+            return grid.view((N, D, H, W, 3))
+        else:
+            base_grid[:, :, :, :, 0] = ivy.linspace(-1, 1, W) * (W - 1) / W
+            base_grid[:, :, :, :, 1] = ivy.expand_dims(
+                ivy.linspace(-1, 1, H) * (H - 1) / H, axis=-1
+            )
+            base_grid[:, :, :, :, 2] = ivy.expand_dims(
+                ivy.expand_dims(ivy.linspace(-1, 1, D) * (D - 1) / D, axis=-1), axis=-1
+            )
+            base_grid[:, :, :, :, 3] = ivy.full((D, H, W), 1)
+            grid = ivy.matmul(base_grid.view((N, D * H * W, 4)), theta.swapaxes(1, 2))
+            return grid.view((N, D, H, W, 3))

{"golden_diff": "diff --git a/ivy/functional/frontends/torch/nn/functional/vision_functions.py b/ivy/functional/frontends/torch/nn/functional/vision_functions.py\n--- a/ivy/functional/frontends/torch/nn/functional/vision_functions.py\n+++ b/ivy/functional/frontends/torch/nn/functional/vision_functions.py\n@@ -319,3 +319,48 @@\n     return interpolate(\n         input, size=size, scale_factor=scale_factor, mode=\"bilinear\", align_corners=True\n     )\n+\n+\n+@with_unsupported_dtypes({\"2.0.1 and below\": (\"float16\", \"bfloat16\")}, \"torch\")\n+@to_ivy_arrays_and_back\n+def affine_grid(theta, size, align_corners=False):\n+    if len(size) == 4:\n+        N, C, H, W = size\n+        base_grid = ivy.empty((N, H, W, 3))\n+        if align_corners:\n+            base_grid[:, :, :, 0] = ivy.linspace(-1, 1, W)\n+            base_grid[:, :, :, 1] = ivy.expand_dims(ivy.linspace(-1, 1, H), axis=-1)\n+            base_grid[:, :, :, 2] = ivy.full((H, W), 1)\n+            grid = ivy.matmul(base_grid.view((N, H * W, 3)), theta.swapaxes(1, 2))\n+            return grid.view((N, H, W, 2))\n+        else:\n+            base_grid[:, :, :, 0] = ivy.linspace(-1, 1, W) * (W - 1) / W\n+            base_grid[:, :, :, 1] = ivy.expand_dims(\n+                ivy.linspace(-1, 1, H) * (H - 1) / H, axis=-1\n+            )\n+            base_grid[:, :, :, 2] = ivy.full((H, W), 1)\n+        grid = ivy.matmul(base_grid.view((N, H * W, 3)), ivy.swapaxes(theta, 1, 2))\n+        return grid.view((N, H, W, 2))\n+    else:\n+        N, C, D, H, W = size\n+        base_grid = ivy.empty((N, D, H, W, 4))\n+        if align_corners:\n+            base_grid[:, :, :, :, 0] = ivy.linspace(-1, 1, W)\n+            base_grid[:, :, :, :, 1] = ivy.expand_dims(ivy.linspace(-1, 1, H), axis=-1)\n+            base_grid[:, :, :, :, 2] = ivy.expand_dims(\n+                ivy.expand_dims(ivy.linspace(-1, 1, D), axis=-1), axis=-1\n+            )\n+            base_grid[:, :, :, :, 3] = ivy.full((D, H, W), 1)\n+            grid = ivy.matmul(base_grid.view((N, D * H * W, 4)), theta.swapaxes(1, 2))\n+            return grid.view((N, D, H, W, 3))\n+        else:\n+            base_grid[:, :, :, :, 0] = ivy.linspace(-1, 1, W) * (W - 1) / W\n+            base_grid[:, :, :, :, 1] = ivy.expand_dims(\n+                ivy.linspace(-1, 1, H) * (H - 1) / H, axis=-1\n+            )\n+            base_grid[:, :, :, :, 2] = ivy.expand_dims(\n+                ivy.expand_dims(ivy.linspace(-1, 1, D) * (D - 1) / D, axis=-1), axis=-1\n+            )\n+            base_grid[:, :, :, :, 3] = ivy.full((D, H, W), 1)\n+            grid = ivy.matmul(base_grid.view((N, D * H * W, 4)), theta.swapaxes(1, 2))\n+            return grid.view((N, D, H, W, 3))\n", "issue": "torch.nn.functional.affine_grid\n\n", "before_files": [{"content": "# global\nimport math\n\n# local\nimport ivy\nfrom ivy import with_unsupported_dtypes\nfrom ivy.functional.frontends.torch.func_wrapper import to_ivy_arrays_and_back\nfrom ivy.utils.exceptions import IvyNotImplementedException\n\n\n@to_ivy_arrays_and_back\ndef pixel_shuffle(input, upscale_factor):\n    input_shape = ivy.shape(input)\n\n    ivy.utils.assertions.check_equal(\n        ivy.get_num_dims(input),\n        4,\n        message=\"pixel_shuffle expects 4D input, but got input with sizes \"\n        + str(input_shape),\n        as_array=False,\n    )\n    b = input_shape[0]\n    c = input_shape[1]\n    h = input_shape[2]\n    w = input_shape[3]\n    upscale_factor_squared = upscale_factor * upscale_factor\n    ivy.utils.assertions.check_equal(\n        c % upscale_factor_squared,\n        0,\n        message=\"pixel_shuffle expects input channel to be divisible by square \"\n        + \"of upscale_factor, but got input with sizes \"\n        + str(input_shape)\n        + \", upscale_factor=\"\n        + str(upscale_factor)\n        + \", and self.size(1)=\"\n        + str(c)\n        + \" is not divisible by \"\n        + str(upscale_factor_squared),\n        as_array=False,\n    )\n    oc = int(c / upscale_factor_squared)\n    oh = h * upscale_factor\n    ow = w * upscale_factor\n\n    input_reshaped = ivy.reshape(input, (b, oc, upscale_factor, upscale_factor, h, w))\n    return ivy.reshape(\n        ivy.permute_dims(input_reshaped, (0, 1, 4, 2, 5, 3)), (b, oc, oh, ow)\n    )\n\n\n@to_ivy_arrays_and_back\ndef pixel_unshuffle(input, downscale_factor):\n    input_shape = ivy.shape(input)\n\n    ivy.utils.assertions.check_equal(\n        ivy.get_num_dims(input),\n        4,\n        message=(\n            f\"pixel_unshuffle expects 4D input, but got input with sizes {input_shape}\"\n        ),\n        as_array=False,\n    ),\n\n    b = input_shape[0]\n    c = input_shape[1]\n    h = input_shape[2]\n    w = input_shape[3]\n    downscale_factor_squared = downscale_factor * downscale_factor\n\n    ivy.utils.assertions.check_equal(\n        [h % downscale_factor, w % downscale_factor],\n        [0, 0],  # Assert h % downscale_factor == 0 and w % downscale_factor == 0\n        message=(\n            \"pixel_unshuffle expects input height and width to be divisible by \"\n            f\"downscale_factor, but got input with sizes {input_shape}\"\n            f\", downscale_factor= {downscale_factor}\"\n            f\", and either self.size(2)= {h}\"\n            f\" or self.size(3)= {w}\"\n            f\" is not divisible by {downscale_factor}\"\n        ),\n        as_array=False,\n    )\n    oc = c * downscale_factor_squared\n    oh = int(h / downscale_factor)\n    ow = int(w / downscale_factor)\n\n    input_reshaped = ivy.reshape(\n        input, (b, c, oh, downscale_factor, ow, downscale_factor)\n    )\n    return ivy.reshape(\n        ivy.permute_dims(input_reshaped, (0, 1, 3, 5, 2, 4)), (b, oc, oh, ow)\n    )\n\n\ndef _handle_padding_shape(padding, n, mode):\n    padding = tuple(\n        [\n            (padding[i * 2], padding[i * 2 + 1])\n            for i in range(int(len(padding) / 2) - 1, -1, -1)\n        ]\n    )\n    while len(padding) < n:\n        if mode == \"circular\":\n            padding = padding + ((0, 0),)\n        else:\n            padding = ((0, 0),) + padding\n    if mode == \"circular\":\n        padding = tuple(list(padding)[::-1])\n    return padding\n\n\n@to_ivy_arrays_and_back\ndef pad(input, pad, mode=\"constant\", value=0):\n    mode_dict = {\n        \"constant\": \"constant\",\n        \"reflect\": \"reflect\",\n        \"replicate\": \"edge\",\n        \"circular\": \"wrap\",\n    }\n    if mode not in mode_dict:\n        raise ValueError(f\"Unsupported padding mode: {mode}\")\n    pad = _handle_padding_shape(pad, len(input.shape), mode)\n    return ivy.pad(input, pad, mode=mode_dict[mode], constant_values=value)\n\n\ndef _get_new_width_height(w_old, h_old, size=None, scale_factor=None):\n    if scale_factor and (not size):\n        if type(scale_factor) == int:\n            h_new = int(w_old * scale_factor)\n            w_new = int(h_old * scale_factor)\n        elif type(scale_factor) == tuple:\n            h_new = int(w_old * scale_factor[0])\n            w_new = int(h_old * scale_factor[1])\n    elif (not scale_factor) and size:\n        if type(size) == int:\n            h_new = size\n            w_new = size\n        elif type(size) == tuple:\n            h_new, w_new = size\n    return h_new, w_new\n\n\n@with_unsupported_dtypes(\n    {\n        \"2.0.1 and below\": (\n            \"bfloat16\",\n            \"float16\",\n        )\n    },\n    \"torch\",\n)\n@to_ivy_arrays_and_back\ndef interpolate(\n    input,\n    size=None,\n    scale_factor=None,\n    mode=\"nearest\",\n    align_corners=None,\n    recompute_scale_factor=None,\n    antialias=False,\n):\n    if mode in [\"nearest\", \"area\", \"nearest-exact\"]:\n        ivy.utils.assertions.check_exists(\n            align_corners,\n            inverse=True,\n            message=(\n                \"align_corners option can only be set with the interpolating modes:\"\n                \" linear | bilinear | bicubic | trilinear\"\n            ),\n        )\n    else:\n        if not ivy.exists(align_corners):\n            align_corners = False\n\n    dim = ivy.get_num_dims(input) - 2  # Number of spatial dimensions.\n\n    if ivy.exists(size) and ivy.exists(scale_factor):\n        raise ivy.utils.exceptions.IvyException(\n            \"only one of size or scale_factor should be defined\"\n        )\n\n    elif ivy.exists(size) and not ivy.exists(scale_factor):\n        scale_factors = None\n\n        if isinstance(size, (list, tuple)):\n            ivy.utils.assertions.check_equal(\n                len(size),\n                dim,\n                inverse=False,\n                message=(\n                    \"Input and output must have the \"\n                    \"same number of spatial dimensions,\"\n                    f\" but got input with spatial dimensions of {list(input.shape[2:])}\"\n                    f\" and output size of {size}. \"\n                    \"Please provide input tensor in (N, C, d1, d2, ...,dK) format\"\n                    \" and output size in (o1, o2, ...,oK) format.\"\n                ),\n                as_array=False,\n            )\n            output_size = size\n        else:\n            output_size = [size for _ in range(dim)]\n\n    elif ivy.exists(scale_factor) and not ivy.exists(size):\n        output_size = None\n\n        if isinstance(scale_factor, (list, tuple)):\n            ivy.utils.assertions.check_equal(\n                len(scale_factor),\n                dim,\n                inverse=False,\n                message=(\n                    \"Input and scale_factor must have the \"\n                    \"same number of spatial dimensions,\"\n                    f\" but got input with spatial dimensions of {list(input.shape[2:])}\"\n                    f\" and scale_factor of shape {scale_factor}. \"\n                    \"Please provide input tensor in (N, C, d1, d2, ...,dK) format\"\n                    \" and scale_factor in (s1, s2, ...,sK) format.\"\n                ),\n                as_array=False,\n            )\n            scale_factors = scale_factor\n        else:\n            scale_factors = [scale_factor for _ in range(dim)]\n\n    else:\n        ivy.utils.assertions.check_any(\n            [ivy.exists(size), ivy.exists(scale_factor)],\n            message=\"either size or scale_factor should be defined\",\n            as_array=False,\n        )\n\n    if (\n        ivy.exists(size)\n        and ivy.exists(recompute_scale_factor)\n        and bool(recompute_scale_factor)\n    ):\n        raise ivy.utils.exceptions.IvyException(\n            \"recompute_scale_factor is not meaningful with an explicit size.\"\n        )\n\n    if ivy.exists(scale_factors):\n        output_size = [\n            math.floor(ivy.shape(input)[i + 2] * scale_factors[i]) for i in range(dim)\n        ]\n\n    if (\n        bool(antialias)\n        and not (mode in [\"bilinear\", \"bicubic\"])\n        and ivy.get_num_dims(input) == 4\n    ):\n        raise ivy.utils.exceptions.IvyException(\n            \"recompute_scale_factor is not meaningful with an explicit size.\"\n        )\n\n    if ivy.get_num_dims(input) == 3 and mode == \"bilinear\":\n        raise IvyNotImplementedException(\n            \"Got 3D input, but bilinear mode needs 4D input\"\n        )\n    if ivy.get_num_dims(input) == 3 and mode == \"trilinear\":\n        raise IvyNotImplementedException(\n            \"Got 3D input, but trilinear mode needs 5D input\"\n        )\n    if ivy.get_num_dims(input) == 4 and mode == \"linear\":\n        raise IvyNotImplementedException(\"Got 4D input, but linear mode needs 3D input\")\n    if ivy.get_num_dims(input) == 4 and mode == \"trilinear\":\n        raise IvyNotImplementedException(\n            \"Got 4D input, but trilinear mode needs 5D input\"\n        )\n    if ivy.get_num_dims(input) == 5 and mode == \"linear\":\n        raise IvyNotImplementedException(\"Got 5D input, but linear mode needs 3D input\")\n    if ivy.get_num_dims(input) == 5 and mode == \"bilinear\":\n        raise IvyNotImplementedException(\n            \"Got 5D input, but bilinear mode needs 4D input\"\n        )\n\n    ivy.utils.assertions.check_elem_in_list(\n        ivy.get_num_dims(input),\n        range(3, 6),\n        message=(\n            \"Input Error: Only 3D, 4D and 5D input Tensors supported (got\"\n            f\" {ivy.get_num_dims(input)}D) for the modes: nearest | linear | bilinear |\"\n            f\" bicubic | trilinear | area | nearest-exact (got {mode})\"\n        ),\n    )\n\n    return ivy.interpolate(\n        input, output_size, mode=mode, align_corners=align_corners, antialias=antialias\n    )\n\n\n@with_unsupported_dtypes({\"2.0.1 and below\": (\"float16\", \"bfloat16\")}, \"torch\")\n@to_ivy_arrays_and_back\ndef upsample(\n    input,\n    size=None,\n    scale_factor=None,\n    mode=\"nearest\",\n    align_corners=None,\n):\n    return interpolate(\n        input,\n        size=size,\n        scale_factor=scale_factor,\n        mode=mode,\n        align_corners=align_corners,\n    )\n\n\n@with_unsupported_dtypes({\"2.0.1 and below\": (\"float16\", \"bfloat16\")}, \"torch\")\n@to_ivy_arrays_and_back\ndef upsample_nearest(input, size=None, scale_factor=None):\n    return interpolate(input, size=size, scale_factor=scale_factor, mode=\"nearest\")\n\n\n@with_unsupported_dtypes({\"2.0.1 and below\": (\"float16\", \"bfloat16\")}, \"torch\")\n@to_ivy_arrays_and_back\ndef upsample_bilinear(input, size=None, scale_factor=None):\n    return interpolate(\n        input, size=size, scale_factor=scale_factor, mode=\"bilinear\", align_corners=True\n    )\n", "path": "ivy/functional/frontends/torch/nn/functional/vision_functions.py"}], "after_files": [{"content": "# global\nimport math\n\n# local\nimport ivy\nfrom ivy import with_unsupported_dtypes\nfrom ivy.functional.frontends.torch.func_wrapper import to_ivy_arrays_and_back\nfrom ivy.utils.exceptions import IvyNotImplementedException\n\n\n@to_ivy_arrays_and_back\ndef pixel_shuffle(input, upscale_factor):\n    input_shape = ivy.shape(input)\n\n    ivy.utils.assertions.check_equal(\n        ivy.get_num_dims(input),\n        4,\n        message=\"pixel_shuffle expects 4D input, but got input with sizes \"\n        + str(input_shape),\n        as_array=False,\n    )\n    b = input_shape[0]\n    c = input_shape[1]\n    h = input_shape[2]\n    w = input_shape[3]\n    upscale_factor_squared = upscale_factor * upscale_factor\n    ivy.utils.assertions.check_equal(\n        c % upscale_factor_squared,\n        0,\n        message=\"pixel_shuffle expects input channel to be divisible by square \"\n        + \"of upscale_factor, but got input with sizes \"\n        + str(input_shape)\n        + \", upscale_factor=\"\n        + str(upscale_factor)\n        + \", and self.size(1)=\"\n        + str(c)\n        + \" is not divisible by \"\n        + str(upscale_factor_squared),\n        as_array=False,\n    )\n    oc = int(c / upscale_factor_squared)\n    oh = h * upscale_factor\n    ow = w * upscale_factor\n\n    input_reshaped = ivy.reshape(input, (b, oc, upscale_factor, upscale_factor, h, w))\n    return ivy.reshape(\n        ivy.permute_dims(input_reshaped, (0, 1, 4, 2, 5, 3)), (b, oc, oh, ow)\n    )\n\n\n@to_ivy_arrays_and_back\ndef pixel_unshuffle(input, downscale_factor):\n    input_shape = ivy.shape(input)\n\n    ivy.utils.assertions.check_equal(\n        ivy.get_num_dims(input),\n        4,\n        message=(\n            f\"pixel_unshuffle expects 4D input, but got input with sizes {input_shape}\"\n        ),\n        as_array=False,\n    ),\n\n    b = input_shape[0]\n    c = input_shape[1]\n    h = input_shape[2]\n    w = input_shape[3]\n    downscale_factor_squared = downscale_factor * downscale_factor\n\n    ivy.utils.assertions.check_equal(\n        [h % downscale_factor, w % downscale_factor],\n        [0, 0],  # Assert h % downscale_factor == 0 and w % downscale_factor == 0\n        message=(\n            \"pixel_unshuffle expects input height and width to be divisible by \"\n            f\"downscale_factor, but got input with sizes {input_shape}\"\n            f\", downscale_factor= {downscale_factor}\"\n            f\", and either self.size(2)= {h}\"\n            f\" or self.size(3)= {w}\"\n            f\" is not divisible by {downscale_factor}\"\n        ),\n        as_array=False,\n    )\n    oc = c * downscale_factor_squared\n    oh = int(h / downscale_factor)\n    ow = int(w / downscale_factor)\n\n    input_reshaped = ivy.reshape(\n        input, (b, c, oh, downscale_factor, ow, downscale_factor)\n    )\n    return ivy.reshape(\n        ivy.permute_dims(input_reshaped, (0, 1, 3, 5, 2, 4)), (b, oc, oh, ow)\n    )\n\n\ndef _handle_padding_shape(padding, n, mode):\n    padding = tuple(\n        [\n            (padding[i * 2], padding[i * 2 + 1])\n            for i in range(int(len(padding) / 2) - 1, -1, -1)\n        ]\n    )\n    while len(padding) < n:\n        if mode == \"circular\":\n            padding = padding + ((0, 0),)\n        else:\n            padding = ((0, 0),) + padding\n    if mode == \"circular\":\n        padding = tuple(list(padding)[::-1])\n    return padding\n\n\n@to_ivy_arrays_and_back\ndef pad(input, pad, mode=\"constant\", value=0):\n    mode_dict = {\n        \"constant\": \"constant\",\n        \"reflect\": \"reflect\",\n        \"replicate\": \"edge\",\n        \"circular\": \"wrap\",\n    }\n    if mode not in mode_dict:\n        raise ValueError(f\"Unsupported padding mode: {mode}\")\n    pad = _handle_padding_shape(pad, len(input.shape), mode)\n    return ivy.pad(input, pad, mode=mode_dict[mode], constant_values=value)\n\n\ndef _get_new_width_height(w_old, h_old, size=None, scale_factor=None):\n    if scale_factor and (not size):\n        if type(scale_factor) == int:\n            h_new = int(w_old * scale_factor)\n            w_new = int(h_old * scale_factor)\n        elif type(scale_factor) == tuple:\n            h_new = int(w_old * scale_factor[0])\n            w_new = int(h_old * scale_factor[1])\n    elif (not scale_factor) and size:\n        if type(size) == int:\n            h_new = size\n            w_new = size\n        elif type(size) == tuple:\n            h_new, w_new = size\n    return h_new, w_new\n\n\n@with_unsupported_dtypes(\n    {\n        \"2.0.1 and below\": (\n            \"bfloat16\",\n            \"float16\",\n        )\n    },\n    \"torch\",\n)\n@to_ivy_arrays_and_back\ndef interpolate(\n    input,\n    size=None,\n    scale_factor=None,\n    mode=\"nearest\",\n    align_corners=None,\n    recompute_scale_factor=None,\n    antialias=False,\n):\n    if mode in [\"nearest\", \"area\", \"nearest-exact\"]:\n        ivy.utils.assertions.check_exists(\n            align_corners,\n            inverse=True,\n            message=(\n                \"align_corners option can only be set with the interpolating modes:\"\n                \" linear | bilinear | bicubic | trilinear\"\n            ),\n        )\n    else:\n        if not ivy.exists(align_corners):\n            align_corners = False\n\n    dim = ivy.get_num_dims(input) - 2  # Number of spatial dimensions.\n\n    if ivy.exists(size) and ivy.exists(scale_factor):\n        raise ivy.utils.exceptions.IvyException(\n            \"only one of size or scale_factor should be defined\"\n        )\n\n    elif ivy.exists(size) and not ivy.exists(scale_factor):\n        scale_factors = None\n\n        if isinstance(size, (list, tuple)):\n            ivy.utils.assertions.check_equal(\n                len(size),\n                dim,\n                inverse=False,\n                message=(\n                    \"Input and output must have the \"\n                    \"same number of spatial dimensions,\"\n                    f\" but got input with spatial dimensions of {list(input.shape[2:])}\"\n                    f\" and output size of {size}. \"\n                    \"Please provide input tensor in (N, C, d1, d2, ...,dK) format\"\n                    \" and output size in (o1, o2, ...,oK) format.\"\n                ),\n                as_array=False,\n            )\n            output_size = size\n        else:\n            output_size = [size for _ in range(dim)]\n\n    elif ivy.exists(scale_factor) and not ivy.exists(size):\n        output_size = None\n\n        if isinstance(scale_factor, (list, tuple)):\n            ivy.utils.assertions.check_equal(\n                len(scale_factor),\n                dim,\n                inverse=False,\n                message=(\n                    \"Input and scale_factor must have the \"\n                    \"same number of spatial dimensions,\"\n                    f\" but got input with spatial dimensions of {list(input.shape[2:])}\"\n                    f\" and scale_factor of shape {scale_factor}. \"\n                    \"Please provide input tensor in (N, C, d1, d2, ...,dK) format\"\n                    \" and scale_factor in (s1, s2, ...,sK) format.\"\n                ),\n                as_array=False,\n            )\n            scale_factors = scale_factor\n        else:\n            scale_factors = [scale_factor for _ in range(dim)]\n\n    else:\n        ivy.utils.assertions.check_any(\n            [ivy.exists(size), ivy.exists(scale_factor)],\n            message=\"either size or scale_factor should be defined\",\n            as_array=False,\n        )\n\n    if (\n        ivy.exists(size)\n        and ivy.exists(recompute_scale_factor)\n        and bool(recompute_scale_factor)\n    ):\n        raise ivy.utils.exceptions.IvyException(\n            \"recompute_scale_factor is not meaningful with an explicit size.\"\n        )\n\n    if ivy.exists(scale_factors):\n        output_size = [\n            math.floor(ivy.shape(input)[i + 2] * scale_factors[i]) for i in range(dim)\n        ]\n\n    if (\n        bool(antialias)\n        and not (mode in [\"bilinear\", \"bicubic\"])\n        and ivy.get_num_dims(input) == 4\n    ):\n        raise ivy.utils.exceptions.IvyException(\n            \"recompute_scale_factor is not meaningful with an explicit size.\"\n        )\n\n    if ivy.get_num_dims(input) == 3 and mode == \"bilinear\":\n        raise IvyNotImplementedException(\n            \"Got 3D input, but bilinear mode needs 4D input\"\n        )\n    if ivy.get_num_dims(input) == 3 and mode == \"trilinear\":\n        raise IvyNotImplementedException(\n            \"Got 3D input, but trilinear mode needs 5D input\"\n        )\n    if ivy.get_num_dims(input) == 4 and mode == \"linear\":\n        raise IvyNotImplementedException(\"Got 4D input, but linear mode needs 3D input\")\n    if ivy.get_num_dims(input) == 4 and mode == \"trilinear\":\n        raise IvyNotImplementedException(\n            \"Got 4D input, but trilinear mode needs 5D input\"\n        )\n    if ivy.get_num_dims(input) == 5 and mode == \"linear\":\n        raise IvyNotImplementedException(\"Got 5D input, but linear mode needs 3D input\")\n    if ivy.get_num_dims(input) == 5 and mode == \"bilinear\":\n        raise IvyNotImplementedException(\n            \"Got 5D input, but bilinear mode needs 4D input\"\n        )\n\n    ivy.utils.assertions.check_elem_in_list(\n        ivy.get_num_dims(input),\n        range(3, 6),\n        message=(\n            \"Input Error: Only 3D, 4D and 5D input Tensors supported (got\"\n            f\" {ivy.get_num_dims(input)}D) for the modes: nearest | linear | bilinear |\"\n            f\" bicubic | trilinear | area | nearest-exact (got {mode})\"\n        ),\n    )\n\n    return ivy.interpolate(\n        input, output_size, mode=mode, align_corners=align_corners, antialias=antialias\n    )\n\n\n@with_unsupported_dtypes({\"2.0.1 and below\": (\"float16\", \"bfloat16\")}, \"torch\")\n@to_ivy_arrays_and_back\ndef upsample(\n    input,\n    size=None,\n    scale_factor=None,\n    mode=\"nearest\",\n    align_corners=None,\n):\n    return interpolate(\n        input,\n        size=size,\n        scale_factor=scale_factor,\n        mode=mode,\n        align_corners=align_corners,\n    )\n\n\n@with_unsupported_dtypes({\"2.0.1 and below\": (\"float16\", \"bfloat16\")}, \"torch\")\n@to_ivy_arrays_and_back\ndef upsample_nearest(input, size=None, scale_factor=None):\n    return interpolate(input, size=size, scale_factor=scale_factor, mode=\"nearest\")\n\n\n@with_unsupported_dtypes({\"2.0.1 and below\": (\"float16\", \"bfloat16\")}, \"torch\")\n@to_ivy_arrays_and_back\ndef upsample_bilinear(input, size=None, scale_factor=None):\n    return interpolate(\n        input, size=size, scale_factor=scale_factor, mode=\"bilinear\", align_corners=True\n    )\n\n\n@with_unsupported_dtypes({\"2.0.1 and below\": (\"float16\", \"bfloat16\")}, \"torch\")\n@to_ivy_arrays_and_back\ndef affine_grid(theta, size, align_corners=False):\n    if len(size) == 4:\n        N, C, H, W = size\n        base_grid = ivy.empty((N, H, W, 3))\n        if align_corners:\n            base_grid[:, :, :, 0] = ivy.linspace(-1, 1, W)\n            base_grid[:, :, :, 1] = ivy.expand_dims(ivy.linspace(-1, 1, H), axis=-1)\n            base_grid[:, :, :, 2] = ivy.full((H, W), 1)\n            grid = ivy.matmul(base_grid.view((N, H * W, 3)), theta.swapaxes(1, 2))\n            return grid.view((N, H, W, 2))\n        else:\n            base_grid[:, :, :, 0] = ivy.linspace(-1, 1, W) * (W - 1) / W\n            base_grid[:, :, :, 1] = ivy.expand_dims(\n                ivy.linspace(-1, 1, H) * (H - 1) / H, axis=-1\n            )\n            base_grid[:, :, :, 2] = ivy.full((H, W), 1)\n        grid = ivy.matmul(base_grid.view((N, H * W, 3)), ivy.swapaxes(theta, 1, 2))\n        return grid.view((N, H, W, 2))\n    else:\n        N, C, D, H, W = size\n        base_grid = ivy.empty((N, D, H, W, 4))\n        if align_corners:\n            base_grid[:, :, :, :, 0] = ivy.linspace(-1, 1, W)\n            base_grid[:, :, :, :, 1] = ivy.expand_dims(ivy.linspace(-1, 1, H), axis=-1)\n            base_grid[:, :, :, :, 2] = ivy.expand_dims(\n                ivy.expand_dims(ivy.linspace(-1, 1, D), axis=-1), axis=-1\n            )\n            base_grid[:, :, :, :, 3] = ivy.full((D, H, W), 1)\n            grid = ivy.matmul(base_grid.view((N, D * H * W, 4)), theta.swapaxes(1, 2))\n            return grid.view((N, D, H, W, 3))\n        else:\n            base_grid[:, :, :, :, 0] = ivy.linspace(-1, 1, W) * (W - 1) / W\n            base_grid[:, :, :, :, 1] = ivy.expand_dims(\n                ivy.linspace(-1, 1, H) * (H - 1) / H, axis=-1\n            )\n            base_grid[:, :, :, :, 2] = ivy.expand_dims(\n                ivy.expand_dims(ivy.linspace(-1, 1, D) * (D - 1) / D, axis=-1), axis=-1\n            )\n            base_grid[:, :, :, :, 3] = ivy.full((D, H, W), 1)\n            grid = ivy.matmul(base_grid.view((N, D * H * W, 4)), theta.swapaxes(1, 2))\n            return grid.view((N, D, H, W, 3))\n", "path": "ivy/functional/frontends/torch/nn/functional/vision_functions.py"}]}