Add scalar support for operators

pavanky · umar456 · commit 7e0fe2776da8 · 2017-07-06T11:55:53.000-04:00
diff --git a/examples/autograd.cpp b/examples/autograd.cpp
@@ -83,6 +83,19 @@ void test_divide_add()
     VERIFY(dy.array() - (1.0 - x.array() / (y.array() * y.array())));
 }
 
+void test_multiply_add_scalar()
+{
+    auto x = Variable(af::randu(5), true);
+    auto y = Variable(af::randu(5), true);
+    auto z = 2 * x + x * y + y;
+    auto dz = Variable(af::constant(1.0, 5), false);
+    z.backward(dz);
+    auto dx = x.grad();
+    auto dy = y.grad();
+    VERIFY(dx.array() - (2.0 + y.array()));
+    VERIFY(dy.array() - (1.0 + x.array()));
+}
+
 int main()
 {
     af::info();
@@ -91,5 +104,6 @@ int main()
     test_no_calc_grad();
     test_multiply_sub();
     test_divide_add();
+    test_multiply_add_scalar();
     return 0;
 }
diff --git a/include/af/autograd/Functions.hpp b/include/af/autograd/Functions.hpp
@@ -18,6 +18,16 @@ namespace af {
         Variable operator -(const Variable &lhs, const Variable &rhs);
         Variable operator /(const Variable &lhs, const Variable &rhs);
 
+        Variable operator +(const double &lhs, const Variable &rhs);
+        Variable operator *(const double &lhs, const Variable &rhs);
+        Variable operator -(const double &lhs, const Variable &rhs);
+        Variable operator /(const double &lhs, const Variable &rhs);
+
+        Variable operator +(const Variable &lhs, const double &rhs);
+        Variable operator *(const Variable &lhs, const double &rhs);
+        Variable operator -(const Variable &lhs, const double &rhs);
+        Variable operator /(const Variable &lhs, const double &rhs);
+
         Variable negate(const Variable &input);
         Variable reciprocal(const Variable &input);
     }
diff --git a/src/autograd/Functions.cpp b/src/autograd/Functions.cpp
@@ -74,5 +74,29 @@ namespace af {
             return Variable(result, {input}, grad_func);
         }
 
+
+#define INSTANTIATE_OPERATOR(OP)                                        \
+        Variable operator OP(const double &lhs_val, const Variable &rhs) \
+        {                                                               \
+            auto lhs = Variable(                                        \
+                af::constant(lhs_val,                                   \
+                             rhs.array().dims(),                        \
+                             rhs.array().type()),                       \
+                false);                                                 \
+            return lhs OP rhs;                                          \
+        }                                                               \
+        Variable operator OP(const Variable &lhs, const double &rhs_val) \
+        {                                                               \
+            auto rhs = Variable(                                        \
+                af::constant(rhs_val,                                   \
+                             lhs.array().dims(), lhs.array().type()),   \
+                false);                                                 \
+            return lhs OP rhs;                                          \
+        }                                                               \
+
+        INSTANTIATE_OPERATOR(+)
+        INSTANTIATE_OPERATOR(-)
+        INSTANTIATE_OPERATOR(*)
+        INSTANTIATE_OPERATOR(/)
     }
 }
