[MPS] Add im2col and col2im to Fallback (#94491)
These are not in the hot path as they are mostly used in Preprocessing layers.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/94491
Approved by: https://github.com/razarmehr
diff --git a/aten/src/ATen/mps/MPSFallback.mm b/aten/src/ATen/mps/MPSFallback.mm
index 69dd47f..822502e 100644
--- a/aten/src/ATen/mps/MPSFallback.mm
+++ b/aten/src/ATen/mps/MPSFallback.mm
@@ -59,6 +59,8 @@
m.impl("repeat_interleave.self_int", torch::CppFunction::makeFromBoxedFunction<&mps_fallback>());
m.impl("_fft_c2c", torch::CppFunction::makeFromBoxedFunction<&mps_fallback>());
m.impl("_fft_r2c", torch::CppFunction::makeFromBoxedFunction<&mps_fallback>());
+ m.impl("im2col", torch::CppFunction::makeFromBoxedFunction<&mps_fallback>()); // Used in preprocessing by nn.Unfold
+ m.impl("col2im", torch::CppFunction::makeFromBoxedFunction<&mps_fallback>());
m.impl("linalg_vector_norm", torch::CppFunction::makeFromBoxedFunction<&mps_fallback>());
m.impl("sgn.out", torch::CppFunction::makeFromBoxedFunction<&mps_fallback>());
m.impl("_slow_conv2d_forward", slow_conv2d_forward_mps);
diff --git a/test/test_mps.py b/test/test_mps.py
index 0d2acca..9e69f65 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -4180,6 +4180,13 @@
helper(3, 1)
+ def test_im2col(self):
+ def helper(x):
+ return torch.nn.functional.unfold(x, kernel_size=(10, 15), dilation=2, padding=5, stride=3)
+ x_cpu = torch.rand(1, 1, 200, 100)
+ x = x_cpu.detach().clone().to('mps')
+ self.assertEqual(helper(x_cpu), helper(x))
+
def test_select(self):
def helper(n, c):
cpu_x = torch.randn(n, c, device='cpu', dtype=torch.float, requires_grad=True)
