Fixed gcd to use 64 bit integers (#21041)

Chillee · facebook-github-bot · commit 2ba608b4a093 · 2019-05-28T16:20:55.000-07:00
Summary: Not much to say. Fixes implementation introduced here: #19115 Pull Request resolved: #21041 Differential Revision: D15528801 Pulled By: Chillee fbshipit-source-id: bacd709eb711ca00156bd70480d6051b437517ed
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -6009,7 +6009,11 @@ def test_gcd(x, y):
             # type: (int, int) -> int
             return math.gcd(x, y)
 
-        for inputs in [(2, 4), (-5, -15), (-5, 15), (10, 0), (0, 10), (-5, 0), (0, -5), (0, 0), (0, -0)]:
+        max_int = 2147483647
+        min_int = -2147483647 - 1
+        int_vals = list(range(-5, 5, 1)) + [max_int + 5, max_int * 2, min_int - 5, min_int * 2]
+        vals = [(i, j) for i in int_vals for j in int_vals]
+        for inputs in vals:
             self.checkScript(test_gcd, inputs)
 
     def test_math_ops1(self):
diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp
@@ -24,8 +24,8 @@
 #include <c10/core/thread_pool.h>
 #include <c10/util/SmallVector.h>
 
-#include <cctype>
 #include <algorithm>
+#include <cctype>
 #include <cmath>
 #include <exception>
 #include <fstream>
@@ -63,8 +63,7 @@ void checkImplicitTensorToNum(at::Tensor t, bool toInt) {
     throw std::runtime_error(
         "Cannot input a tensor of dimension other than 0 as a scalar argument");
   }
-  if (toInt &&
-      !isIntegralType(t.scalar_type())) {
+  if (toInt && !isIntegralType(t.scalar_type())) {
     std::stringstream ss;
     ss << "Cannot input a tensor of type " << t.scalar_type()
        << " as an integral argument";
@@ -98,9 +97,9 @@ static int64_t floordiv(int64_t a, int64_t b) {
   }
 }
 
-static int gcd(int a, int b) {
+static int64_t gcd(int64_t a, int64_t b) {
   while (b != 0) {
-    int r = a % b;
+    int64_t r = a % b;
     a = b;
     b = r;
   }
