Split should return TransformedTensors

Author: Atry

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

@@ -937,10 +937,10 @@ trait Tensors extends OpenCL {
     /**
       * @group delayed
       */
-    def split(dimension: Int): IndexedSeq[Tensor] = {
+    def split(dimension: Int): IndexedSeq[TransformedTensor] = {
       // TODO: override map/reduce to produce less OpenCL C code
       val newShape = shape.patch(dimension, Nil, 1)
-      new IndexedSeq[Tensor] {
+      new IndexedSeq[TransformedTensor] {
 
         val length: Int = shape(dimension)
 

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

@@ -307,7 +307,7 @@ class TensorsSpec extends AsyncFreeSpec with Matchers {
         val Array(`j`, k) = matrix2.shape
         val product = matrix1.broadcast(Array(i, j, k)) * matrix2.reshape(Array(1, j, k)).broadcast(Array(i, j, k))
 
-        product.split(1).reduce(_ + _)
+        product.split(1).reduce[Tensor](_ + _)
 
       }
 

benchmarks/src/jmh/scala/com/thoughtworks/compute/benchmarks.scala

@@ -100,7 +100,7 @@ object benchmarks {
           // unroll only j
           val Array(`j`, k) = matrix2.shape
           val product = matrix1.broadcast(Array(i, j, k)) * matrix2.reshape(Array(1, j, k)).broadcast(Array(i, j, k))
-          product.split(1).reduce(_ + _)
+          product.split(1).reduce[Tensor](_ + _)
         }
       }