From fe81dea243860ad063d5490d13225182ddbfaedc Mon Sep 17 00:00:00 2001
From: Andrei Kislitsyn <kisandy@yandex.ru>
Date: Sat, 1 May 2021 14:22:05 +0300
Subject: [PATCH] stack

---
 .../core/algebras/DoubleTensorAlgebra.kt      | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleTensorAlgebra.kt b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleTensorAlgebra.kt
index f428b9d2e..c6fb301b5 100644
--- a/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleTensorAlgebra.kt
+++ b/kmath-tensors/src/commonMain/kotlin/space/kscience/kmath/tensors/core/algebras/DoubleTensorAlgebra.kt
@@ -382,7 +382,11 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra<Double> {
             return false
         }
         for (i in 0 until n) {
-            if (!eqFunction(tensor.mutableBuffer[tensor.bufferStart + i], other.tensor.mutableBuffer[other.tensor.bufferStart + i])) {
+            if (!eqFunction(
+                    tensor.mutableBuffer[tensor.bufferStart + i],
+                    other.tensor.mutableBuffer[other.tensor.bufferStart + i]
+                )
+            ) {
                 return false
             }
         }
@@ -395,4 +399,20 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra<Double> {
     public fun TensorStructure<Double>.randNormalLike(seed: Long = 0): DoubleTensor =
         DoubleTensor(tensor.shape, getRandomNormals(tensor.shape.reduce(Int::times), seed))
 
+    // stack tensors by axis 0
+    public fun stack(tensors: List<DoubleTensor>): DoubleTensor {
+        val shape = tensors.firstOrNull()?.shape
+        check(shape != null) { "Collection must have at least 1 element" }
+        check(tensors.all { it.shape contentEquals shape }) {"Stacking tensors must have same shapes"}
+        val resShape = intArrayOf(tensors.size) + shape
+        val resBuffer = tensors.flatMap {
+            it.tensor.mutableBuffer.array().drop(it.bufferStart).take(it.numElements)
+        }.toDoubleArray()
+        return DoubleTensor(resShape, resBuffer, 0)
+    }
+
+    // build tensor from this rows by given indices
+    public fun TensorStructure<Double>.rowsByIndices(indices: IntArray): DoubleTensor {
+        return stack(indices.map { this[it] })
+    }
 }