log_softmax: fix meta function output argument dtype check. (pytorch#140289)

Tracking issue: pytorch#138399
Pull Request resolved: pytorch#140289
Approved by: https://github.com/ezyang
ghstack dependencies: pytorch#140186, pytorch#140286, pytorch#140288

Author: ysiraichi

test/test_ops.py

@@ -168,7 +168,6 @@ def reduction_dtype_filter(op):
     xfail("linalg.solve"),
     xfail("linalg.solve_ex"),
     xfail("linalg.solve_triangular"),
-    xfail("log_softmax"),
     xfail("logcumsumexp"),
     xfail("lu_solve"),
     xfail("lu_unpack"),

torch/_decomp/decompositions.py

@@ -1220,7 +1220,7 @@ def _softmax(x: Tensor, dim: int, half_to_float: bool):
 
 
 @register_decomposition(aten._log_softmax)
-@out_wrapper()
+@out_wrapper(exact_dtype=True)
 def _log_softmax(x: Tensor, dim: int, half_to_float: bool):
     # eager log_softmax returns a contiguous tensor. Ensure that decomp also
     # returns a contiguous tensor.