add test case for executor reshaping

liyangdev · Jun 5, 2016 · a76ab7c · a76ab7c
1 parent 1cfd6d6
commit a76ab7c
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/scala-package/core/src/test/scala/ml/dmlc/mxnet/ExecutorSuite.scala b/scala-package/core/src/test/scala/ml/dmlc/mxnet/ExecutorSuite.scala
@@ -40,4 +40,24 @@ class ExecutorSuite extends FunSuite with BeforeAndAfterAll {
     assert(reldiff(lhsGrad, lhsGrad2) < 1e-6)
     assert(reldiff(rhsGrad, rhsGrad2) < 1e-6)
   }
+
+  test("reshape") {
+    val x = Symbol.Variable("x")
+    val y = Symbol.FullyConnected()(Map("data" -> x, "num_hidden" -> 4))
+
+    val exec = y.simpleBind(Context.cpu(), "write", shapeDict = Map("x" -> Shape(5, 4)))
+    exec.argArrays(0).set(1)
+    exec.argArrays(1).set(1)
+    exec.argArrays(2).set(0)
+
+    val newExec = exec.reshape(kwargs = Map("x" -> Shape(3, 4)))
+    newExec.forward(isTrain = false)
+    // test sub exec forward
+    assert(newExec.outputs(0).toArray.forall(_ == 4))
+    // test shared memory
+    assert(exec.outputs(0).toArray.take(3).forall(_ == 4))
+    // test base exec forward
+    exec.forward(isTrain = false)
+    assert(exec.outputs(0).toArray.forall(_ == 4))
+  }
 }