Rewrite some decomps to allow producing aten ops (#93099)

SherlockNoMad · pytorchmergebot · commit 438f12d91a8a · 2023-01-31T17:46:20.000Z
This introduces a new stop to the decomposition train. Before reaching prims.view_of, it will stop at aten.alias. Export path wants to get off the train at aten ops. Pull Request resolved: #93099 Approved by: https://github.com/ngimel
diff --git a/test/expect/HasDecompTest.test_has_decomposition.expect b/test/expect/HasDecompTest.test_has_decomposition.expect
@@ -559,7 +559,6 @@ aten::addmv.out
 aten::addr_
 aten::affine_grid_generator
 aten::affine_grid_generator.out
-aten::alias
 aten::alias_copy
 aten::alias_copy.out
 aten::allclose
diff --git a/torch/_prims_common/wrappers.py b/torch/_prims_common/wrappers.py
@@ -34,12 +34,9 @@ def _maybe_convert_to_dtype(a: None, dtype: torch.dtype) -> None:
 
 # TODO: implement ref.cast with an option to enforce safe casting
 def _maybe_convert_to_dtype(a, dtype):
-    import torch._prims as prims
     if isinstance(a, TensorLike):
         if a.dtype != dtype:
-            # NOTE: this is incorrect on the CPU
-            # See https://github.com/pytorch/pytorch/issues/77553
-            return prims.convert_element_type(a, dtype)
+            return a.to(dtype)
         return a
     if isinstance(a, Number):
         return utils.dtype_to_type_ctor(dtype)(a)  # type: ignore[arg-type]
diff --git a/torch/_refs/__init__.py b/torch/_refs/__init__.py
@@ -3905,12 +3905,17 @@ def T(a: TensorLikeType) -> TensorLikeType:
     return a.t()
 
 
+@register_decomposition(aten.alias)
+def alias(a: TensorLikeType) -> TensorLikeType:
+    return prims.view_of(a)
+
+
 @register_decomposition(aten.transpose)
 def transpose(a: TensorLikeType, dim0: int, dim1: int) -> TensorLikeType:
     _dim0, _dim1 = utils.canonicalize_dims(a.ndim, (dim0, dim1))  # type: ignore[misc]
 
     if a.ndim <= 1 or dim0 == dim1:
-        return prims.view_of(a)
+        return aten.alias.default(a)
 
     _permutation = list(range(0, a.ndim))
     _permutation[_dim0] = _dim1
