Fix broadcast method

Author: Atry

OpenCLKernelBuilder/src/main/scala/com/thoughtworks/compute/OpenCLKernelBuilder.scala

@@ -298,7 +298,7 @@ trait OpenCLKernelBuilder extends AllExpressions {
     def extract: Element = {
       val numberOfRows = originalShape.length
       val numberOfColumns = matrix.length / numberOfRows
-      if (matrix.length % numberOfRows != 0) {
+      if (matrix.length != numberOfRows * numberOfColumns) {
         throw new IllegalStateException()
       }
 

Tensors/src/main/scala/com/thoughtworks/compute/Tensors.scala

@@ -711,7 +711,7 @@ trait Tensors extends OpenCL {
         if (i < length) {
           shape(i) match {
             case di if di == newShape(i) =>
-              matrix1(i * (length + 1) + i) = 1.0
+              matrix1(i * (newLength + 1) + i) = 1.0
             case 1 =>
             case _ =>
               throw new IllegalArgumentException(

Tensors/src/test/scala/com/thoughtworks/compute/TensorsSpec.scala

@@ -293,4 +293,14 @@ class TensorsSpec extends AsyncFreeSpec with Matchers {
     .run
     .toScalaFuture
 
+  "broadcast" in doTensors
+    .map { tensors =>
+      import tensors._
+
+      val matrix1 = Tensor(Array(Array(1.0f, 2.0f, 3.0f), Array(4.0f, 5.0f, 6.0f)))
+      matrix1.broadcast(Array(2, 3, 4)).toString should be(
+        "[[[1.0,1.0,1.0,1.0],[2.0,2.0,2.0,2.0],[3.0,3.0,3.0,3.0]],[[4.0,4.0,4.0,4.0],[5.0,5.0,5.0,5.0],[6.0,6.0,6.0,6.0]]]")
+    }
+    .run
+    .toScalaFuture
 }