[pytorch][tensorexpr] Promote integer arguments to sin/cos/tan to flo…

…at (#46776)

Summary:
Pull Request resolved: #46776

Following numpy and (now) eager mode

Fixes #46458

Test Plan: test_jit_fuser_te

Reviewed By: navahgar

Differential Revision: D24509884

fbshipit-source-id: c063030fc609ba4aefcd9abd25b50f082fef1548

test/test_jit_fuser_te.py

@@ -1274,7 +1274,7 @@ def apply(fn):
             torch.erf,
             torch.erfc,
             torch.cos,
-            # torch.sin, Note: Reference https://github.com/pytorch/pytorch/issues/46458
+            torch.sin,
             torch.tan,
             torch.acos,
             torch.asin,

torch/csrc/jit/tensorexpr/kernel.cpp

@@ -411,6 +411,16 @@ ExprHandle TensorExprKernel::constant(const torch::jit::Value* v) {
   return scalars_.at(v->unique());
 }
 
+ExprHandle promoteIntegerToFloat(const ExprHandle& e) {
+  auto scalarType = static_cast<c10::ScalarType>(e.dtype().scalar_type());
+  if (!c10::isIntegralType(scalarType)) {
+    return e;
+  }
+  auto defaultType = static_cast<tensorexpr::ScalarType>(
+      c10::typeMetaToScalarType(c10::get_default_dtype()));
+  return Cast::make(Dtype(defaultType, e.dtype().lanes()), e);
+}
+
 void TensorExprKernel::promoteInputs(std::vector<ExprHandle>& inputs) {
   if (inputs.empty()) {
     return;
@@ -965,18 +975,21 @@ Tensor* TensorExprKernel::computeValue(const torch::jit::Value* v) {
     } break;
 
     case aten::cos: {
-      return computeOneOperand(
-          "aten_cos", v, [](const ExprHandle& a) { return cos(a); });
+      return computeOneOperand("aten_cos", v, [](const ExprHandle& a) {
+        return cos(promoteIntegerToFloat(a));
+      });
     } break;
 
     case aten::sin: {
-      return computeOneOperand(
-          "aten_sin", v, [](const ExprHandle& a) { return sin(a); });
+      return computeOneOperand("aten_sin", v, [](const ExprHandle& a) {
+        return sin(promoteIntegerToFloat(a));
+      });
     } break;
 
     case aten::tan: {
-      return computeOneOperand(
-          "aten_tan", v, [](const ExprHandle& a) { return tan(a); });
+      return computeOneOperand("aten_tan", v, [](const ExprHandle& a) {
+        return tan(promoteIntegerToFloat(a));
+      });
     } break;
 
     case aten::type_as: {