From b7cac3a015d235d81d02d3bb4b8834dda2b48b95 Mon Sep 17 00:00:00 2001
From: Andrei Kislitsyn <kisandy@yandex.ru>
Date: Sat, 1 May 2021 13:32:50 +0300
Subject: [PATCH] fix sequences + array casting

---
 .../kmath/tensors/core/BufferedTensor.kt      | 38 ++++++++++++++++++-
 .../kscience/kmath/tensors/core/linUtils.kt   |  4 +-
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/BufferedTensor.kt b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/BufferedTensor.kt
index 4e286b489..867b4fb7a 100644
--- a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/BufferedTensor.kt
+++ b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/BufferedTensor.kt
@@ -49,6 +49,7 @@ public class DoubleTensor internal constructor(
 
 internal fun BufferedTensor<Int>.asTensor(): IntTensor =
     IntTensor(this.shape, this.mutableBuffer.array(), this.bufferStart)
+
 internal fun BufferedTensor<Double>.asTensor(): DoubleTensor =
     DoubleTensor(this.shape, this.mutableBuffer.array(), this.bufferStart)
 
@@ -78,4 +79,39 @@ internal val TensorStructure<Int>.tensor: IntTensor
     }
 
 public fun TensorStructure<Double>.toDoubleTensor(): DoubleTensor = this.tensor
-public fun TensorStructure<Int>.toIntTensor(): IntTensor = this.tensor
\ No newline at end of file
+public fun TensorStructure<Int>.toIntTensor(): IntTensor = this.tensor
+
+public fun Array<DoubleArray>.toDoubleTensor(): DoubleTensor {
+    val n = size
+    check(n > 0) { "An empty array cannot be casted to tensor" }
+    val m = first().size
+    check(m > 0) { "Inner arrays must have at least 1 argument" }
+    check(all { size == m }) { "Inner arrays must be the same size" }
+
+    val shape = intArrayOf(n, m)
+    val buffer = this.flatMap { arr -> arr.map { it } }.toDoubleArray()
+
+    return DoubleTensor(shape, buffer, 0)
+}
+
+
+public fun Array<IntArray>.toIntTensor(): IntTensor {
+    val n = size
+    check(n > 0) { "An empty array cannot be casted to tensor" }
+    val m = first().size
+    check(m > 0) { "Inner arrays must have at least 1 argument" }
+    check(all { size == m }) { "Inner arrays must be the same size" }
+
+    val shape = intArrayOf(n, m)
+    val buffer = this.flatMap { arr -> arr.map { it } }.toIntArray()
+
+    return IntTensor(shape, buffer, 0)
+}
+
+public fun DoubleTensor.toDoubleArray(): DoubleArray {
+    return tensor.mutableBuffer.array().drop(bufferStart).take(numElements).toDoubleArray()
+}
+
+public fun IntTensor.toIntArray(): IntArray {
+    return tensor.mutableBuffer.array().drop(bufferStart).take(numElements).toIntArray()
+}
\ No newline at end of file
diff --git a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/linUtils.kt b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/linUtils.kt
index a152b3a17..e54cc4d26 100644
--- a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/linUtils.kt
+++ b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/linUtils.kt
@@ -18,7 +18,7 @@ internal fun <T> BufferedTensor<T>.vectorSequence(): Sequence<BufferedTensor<T>>
     val vectorOffset = shape[n - 1]
     val vectorShape = intArrayOf(shape.last())
     for (offset in 0 until numElements step vectorOffset) {
-        val vector = BufferedTensor(vectorShape, mutableBuffer, offset)
+        val vector = BufferedTensor(vectorShape, mutableBuffer, bufferStart + offset)
         yield(vector)
     }
 }
@@ -29,7 +29,7 @@ internal fun <T> BufferedTensor<T>.matrixSequence(): Sequence<BufferedTensor<T>>
     val matrixOffset = shape[n - 1] * shape[n - 2]
     val matrixShape = intArrayOf(shape[n - 2], shape[n - 1])
     for (offset in 0 until numElements step matrixOffset) {
-        val matrix = BufferedTensor(matrixShape, mutableBuffer, offset)
+        val matrix = BufferedTensor(matrixShape, mutableBuffer, bufferStart + offset)
         yield(matrix)
     }
 }