Add pixel_shuffle to core aten decomps (#119899)

angelayi · pytorchmergebot · commit 9201d7335a25 · 2024-02-14T21:01:11.000Z
Summary: #118239 added a decomposition for pixel_shuffle, so pixel_shuffle no longer needs to be a Core ATen Op. We have also fixed the internal use case so that it no longer special cases on pixel_shuffle, allowing us to revert the changes in #118921. Test Plan: CI Differential Revision: D53766709 Pull Request resolved: #119899 Approved by: https://github.com/peterbell10, https://github.com/lezcano
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -4463,7 +4463,6 @@
     MPS: pixel_shuffle_mps
     CompositeExplicitAutogradNonFunctional: math_pixel_shuffle
   autogen: pixel_shuffle.out
-  tags: core
 
 - func: pixel_unshuffle(Tensor self, int downscale_factor) -> Tensor
   dispatch:
diff --git a/test/expect/HasDecompTest.test_aten_core_operators.expect b/test/expect/HasDecompTest.test_aten_core_operators.expect
@@ -388,8 +388,6 @@ aten::normal.float_float_out
 aten::normal.out
 aten::normal_
 aten::permute
-aten::pixel_shuffle
-aten::pixel_shuffle.out
 aten::polar
 aten::polar.out
 aten::pow.Scalar
diff --git a/torch/_decomp/__init__.py b/torch/_decomp/__init__.py
@@ -371,6 +371,7 @@ def core_aten_decompositions() -> Dict[torch._ops.OperatorBase, Callable]:
             aten.norm,
             aten.ones,
             aten.ones_like,
+            aten.pixel_shuffle,
             aten.pixel_unshuffle,
             aten._prelu_kernel,
             aten._prelu_kernel_backward,
diff --git a/torch/_inductor/decomposition.py b/torch/_inductor/decomposition.py
@@ -57,7 +57,6 @@
         aten.native_batch_norm,
         aten.native_group_norm,
         aten.native_layer_norm,
-        aten.pixel_shuffle,
         aten._softmax,
         aten.sin_,
         aten.sqrt_,
