Adding expandAs, reduceAs, and transpose

pavanky · umar456 · commit 546a87a436bf · 2017-07-06T11:55:53.000-04:00
diff --git a/examples/autograd.cpp b/examples/autograd.cpp
@@ -128,6 +128,32 @@ void test_tanh()
     VERIFY(dx.array() - (1 + af::tanh(x.array())) * (1 - af::tanh(x.array())));
 }
 
+void test_expand()
+{
+    auto x = Variable(af::randu(5), true);
+    auto y = Variable(af::randu(5, 2), true);
+    auto z = y * expandAs(x, y);
+    auto dz = Variable(af::constant(1.0, 5, 2), false);
+    z.backward(dz);
+    auto dy = y.grad();
+    auto dx = x.grad();
+    VERIFY(dy.array() - af::tile(x.array(), 1, 2));
+    VERIFY(dx.array() - af::sum(y.array(), 1));
+}
+
+void test_reduce()
+{
+    auto x = Variable(af::randu(5), true);
+    auto y = Variable(af::randu(5, 2), true);
+    auto z = x * reduceAs(y, x);
+    auto dz = Variable(af::constant(1.0, 5), false);
+    z.backward(dz);
+    auto dy = y.grad();
+    auto dx = x.grad();
+    VERIFY(dy.array() - af::tile(x.array(), 1, 2));
+    VERIFY(dx.array() - af::sum(y.array(), 1));
+}
+
 int main()
 {
     af::info();
@@ -140,5 +166,7 @@ int main()
     test_exp();
     test_sigmoid();
     test_tanh();
+    test_expand();
+    test_reduce();
     return 0;
 }
diff --git a/include/af/autograd/Functions.hpp b/include/af/autograd/Functions.hpp
@@ -36,5 +36,9 @@ namespace af {
         Variable cos(const Variable &input);
         Variable tanh(const Variable &input);
         Variable sigmoid(const Variable &input);
+
+        Variable transpose(const Variable &input);
+        Variable expandAs(const Variable &input, const Variable &reference);
+        Variable reduceAs(const Variable &input, const Variable &reference);
     }
 }
diff --git a/src/autograd/Functions.cpp b/src/autograd/Functions.cpp
@@ -147,5 +147,42 @@ namespace af {
             return Variable(result, {input}, grad_func);
         }
 
+        Variable transpose(const Variable &input)
+        {
+            auto result = transpose(input.array());
+            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
+                inputs[0].addGrad(transpose(grad_output));
+            };
+            return Variable(result, {input}, grad_func);
+        }
+
+        Variable expandAs(const Variable &input, const Variable &reference)
+        {
+            dim4 dims(1,1,1,1);
+            dim4 idims = input.array().dims();
+            dim4 rdims = reference.array().dims();
+            for (int i = 0; i < 4; i++) {
+                dims[i] = rdims[i] / idims[i];
+            }
+            auto result = tile(input.array(), dims);
+            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
+                inputs[0].addGrad(reduceAs(grad_output, inputs[0]));
+            };
+            return Variable(result, {input}, grad_func);
+        }
+
+        Variable reduceAs(const Variable &input, const Variable &reference)
+        {
+            dim4 idims = input.array().dims();
+            dim4 rdims = reference.array().dims();
+            auto result = input.array();
+            for (int i = 0; i < 4; i++) {
+                if (idims[i] != rdims[i]) result = sum(result, i);
+            }
+            auto grad_func = [](std::vector<Variable> &inputs, const Variable &grad_output) {
+                inputs[0].addGrad(expandAs(grad_output, inputs[0]));
+            };
+            return Variable(result, {input}, grad_func);
+        }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -36,5 +36,9 @@ namespace af {`
`36`	`36`	`Variable cos(const Variable &input);`
`37`	`37`	`Variable tanh(const Variable &input);`
`38`	`38`	`Variable sigmoid(const Variable &input);`
	`39`	`+`
	`40`	`+ Variable transpose(const Variable &input);`
	`41`	`+ Variable expandAs(const Variable &input, const Variable &reference);`
	`42`	`+ Variable reduceAs(const Variable &input, const Variable &reference);`
`39`	`43`	`}`
`40`	`44`	`}`