Added base parameter to math.log (#21151)

Chillee · facebook-github-bot · commit 80020306ef55 · 2019-05-30T13:32:52.000-07:00
Summary: Pull Request resolved: #21151 ghimport-source-id: 76dc085 Differential Revision: D15563185 Pulled By: Chillee fbshipit-source-id: 6ed7cc32ed7c103f360022b97f6df47ccd0403e7
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -5970,6 +5970,10 @@ def test_log_float(x):
             # type: (float) -> float
             return math.log(x)
 
+        def test_log_base_float(x, y):
+            # type: (float, float) -> float
+            return math.log(x, y)
+
         def test_log1p_int(x):
             # type: (int) -> float
             return math.log1p(x)
@@ -6006,6 +6010,7 @@ def test_pow_float(x, y):
             # type: (float, float) -> float
             return math.pow(x, y)
 
+
         def test_pow_int(x, y):
             # type: (float, int) -> float
             return math.pow(x, y)
@@ -6014,6 +6019,7 @@ def test_pow_int(x, y):
         self.checkScript(test_ceil, (1.5,))
         self.checkScript(test_log_int, (2,))
         self.checkScript(test_log_float, (2.0,))
+        self.checkScript(test_log_base_float, (2.0, 5.0))
         self.checkScript(test_log1p_int, (1,))
         self.checkScript(test_log1p_float, (1.0,))
         self.checkScript(test_log10_int, (2,))
@@ -6023,6 +6029,7 @@ def test_pow_int(x, y):
         self.checkScript(test_sqrt_int, (2,))
         self.checkScript(test_sqrt_float, (2.0,))
         self.checkScript(test_pow_float, (2.0, 2.0))
+        self.checkScript(test_pow_float, (2.0, 2.0))
         self.checkScript(test_pow_int, (2.0, 2))
 
     @unittest.skipIf(PY2, "Requires python 3")
diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp
@@ -1078,23 +1078,30 @@ RegisterOperators logging_operators(
 #define DEFINE_BINARY_OP(aten_op, op)             \
   DEFINE_GENERIC_OP(aten_op, op, op, int, float), \
       DEFINE_INT_FLOAT_OP(aten_op, op, float)
+
+#define DEFINE_BINARY_FLOAT_OP(aten_op, op)         \
+  DEFINE_GENERIC_OP(aten_op, op, op, float, float), \
+      DEFINE_INT_FLOAT_OP(aten_op, op, float)
+
 #define DEFINE_COMPARISON_OP(aten_op, op)         \
   DEFINE_GENERIC_OP(aten_op, op, op, bool, bool), \
       DEFINE_INT_FLOAT_OP(aten_op, op, bool), DEFINE_STR_CMP_OP(aten_op, op)
 
 #define DEFINE_UNARY_OP(aten_op, op, int_result, float_result)            \
-  Operator(#aten_op "(int a) -> " #int_result, [](Stack& stack) {         \
-    int64_t a;                                                            \
-    pop(stack, a);                                                        \
-    push(stack, op);                                                      \
-    return 0;                                                             \
-  }),                                                                     \
-  Operator(#aten_op "(float a) -> " #float_result, [](Stack& stack) {     \
-    double a;                                                             \
-    pop(stack, a);                                                        \
-    push(stack, op);                                                      \
-    return 0;                                                             \
-  })
+  Operator(                                                               \
+      #aten_op "(int a) -> " #int_result,                                 \
+      [](Stack& stack) {                                                  \
+        int64_t a;                                                        \
+        pop(stack, a);                                                    \
+        push(stack, op);                                                  \
+        return 0;                                                         \
+      }),                                                                 \
+      Operator(#aten_op "(float a) -> " #float_result, [](Stack& stack) { \
+        double a;                                                         \
+        pop(stack, a);                                                    \
+        push(stack, op);                                                  \
+        return 0;                                                         \
+      })
 
 #define DEFINE_BOOL_OP(aten_op, op)                                \
   Operator(#aten_op "(bool a, bool b) -> bool", [](Stack& stack) { \
@@ -2124,6 +2131,7 @@ RegisterOperators reg2({
     DEFINE_UNARY_OP(aten::floor, std::floor(a), float, float),
     DEFINE_UNARY_OP(aten::ceil, std::ceil(a), float, float),
     DEFINE_UNARY_OP(aten::log, std::log(a), float, float),
+    DEFINE_BINARY_FLOAT_OP(aten::log, std::log(a) / std::log(b)),
     DEFINE_UNARY_OP(aten::log1p, std::log1p(a), float, float),
     DEFINE_UNARY_OP(aten::log10, std::log10(a), float, float),
     DEFINE_UNARY_OP(aten::exp, std::exp(a), float, float),
