Repository: ivy-llc/ivy

Issue:
torch.nn.functional.affine_grid



Discussion:

Resolving patch:
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 @@ def upsample_bilinear(input, size=None, scale_factor=None):
     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))


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_torch/test_vision_functions.py b/ivy_tests/test_ivy/test_frontends/test_torch/test_vision_functions.py
--- a/ivy_tests/test_ivy/test_frontends/test_torch/test_vision_functions.py
+++ b/ivy_tests/test_ivy/test_frontends/test_torch/test_vision_functions.py
@@ -274,3 +274,72 @@ def test_torch_upsample_bilinear(
         input=x[0],
         size=size,
     )
+
+
+@st.composite
+def _affine_grid_helper(draw):
+    align_corners = draw(st.booleans())
+    dims = draw(st.integers(4, 5))
+    if dims == 4:
+        size = draw(
+            st.tuples(
+                st.integers(1, 20),
+                st.integers(1, 20),
+                st.integers(2, 20),
+                st.integers(2, 20),
+            )
+        )
+        theta_dtype, theta = draw(
+            helpers.dtype_and_values(
+                available_dtypes=helpers.get_dtypes("float"),
+                min_value=0,
+                max_value=1,
+                shape=(size[0], 2, 3),
+            )
+        )
+        return theta_dtype, theta[0], size, align_corners
+    else:
+        size = draw(
+            st.tuples(
+                st.integers(1, 20),
+                st.integers(1, 20),
+                st.integers(1, 20),
+                st.integers(2, 20),
+                st.integers(2, 20),
+            )
+        )
+        theta_dtype, theta = draw(
+            helpers.dtype_and_values(
+                available_dtypes=helpers.get_dtypes("float"),
+                min_value=0,
+                max_value=1,
+                shape=(size[0], 3, 4),
+            )
+        )
+        return theta_dtype, theta[0], size, align_corners
+
+
+@handle_frontend_test(
+    fn_tree="torch.nn.functional.affine_grid",
+    dtype_and_input_and_other=_affine_grid_helper(),
+)
+def test_torch_affine_grid(
+    *,
+    dtype_and_input_and_other,
+    on_device,
+    fn_tree,
+    frontend,
+    test_flags,
+):
+    dtype, theta, size, align_corners = dtype_and_input_and_other
+
+    helpers.test_frontend_function(
+        input_dtypes=dtype,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        theta=theta,
+        size=size,
+        align_corners=align_corners,
+    )