[MPS] Enable adaptive avg pool 2d with larger output size (#85726)
* Handle adpative pool 2d forward and backward when ouptut size is larger than input size
* Disallow larger output size if not a multiple of input size
Fixes: https://github.com/pytorch/pytorch/issues/80732
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85726
Approved by: https://github.com/malfet
diff --git a/test/test_mps.py b/test/test_mps.py
index 19de49f..635acd5 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -3861,6 +3861,21 @@
helper((2, 16, 16), (4, 4), False)
+ # Output shape larger than input shape
+
+ helper((2, 2, 4, 4), (8, 8), False)
+ helper((2, 2, 2, 2), (4, 4), False)
+ helper((2, 2, 3, 3), (9, 9), False)
+ helper((2, 2, 2, 2), (16, 16), False)
+ helper((2, 2, 2, 16), (16, 16), False)
+
+ helper((2, 4, 4), (16, 16), False)
+
+ try:
+ helper((2, 2, 3, 3), (7, 7), False)
+ except Exception as e:
+ pass
+
# Test max avg pool2d - when the input size is a multiple of output size
# Not testing for channels last right now
def test_adaptive_max_pool2d_simple(self):
