From 384415dc98ae27125f894798057f6052103e2594 Mon Sep 17 00:00:00 2001 From: Roland Grinis Date: Sat, 13 Mar 2021 19:16:13 +0000 Subject: [PATCH] utils module for tensors --- .../kscience/kmath/tensors/BufferedTensor.kt | 28 ++++-- .../kmath/tensors/RealTensorAlgebra.kt | 63 ------------ .../kscience/kmath/tensors/TensorAlgebra.kt | 46 +-------- .../kscience/kmath/tensors/TensorStrides.kt | 6 +- .../space/kscience/kmath/tensors/utils.kt | 96 +++++++++++++++++++ .../kmath/tensors/TestRealTensorAlgebra.kt | 8 +- 6 files changed, 124 insertions(+), 123 deletions(-) create mode 100644 kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/utils.kt diff --git a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/BufferedTensor.kt b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/BufferedTensor.kt index 9ffe2db61..29605024d 100644 --- a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/BufferedTensor.kt +++ b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/BufferedTensor.kt @@ -1,9 +1,8 @@ package space.kscience.kmath.tensors -import space.kscience.kmath.linear.BufferMatrix import space.kscience.kmath.nd.MutableNDBuffer import space.kscience.kmath.structures.* -import space.kscience.kmath.structures.BufferAccessor2D + public open class BufferedTensor( override val shape: IntArray, @@ -15,13 +14,13 @@ public open class BufferedTensor( buffer ) { - public operator fun get(i: Int, j: Int): T{ - check(this.dimension == 2) {"Not matrix"} - return this[intArrayOf(i, j)] + public operator fun get(i: Int, j: Int): T { + check(this.dimension == 2) { "Not matrix" } + return this[intArrayOf(i, j)] } - public operator fun set(i: Int, j: Int, value: T): Unit{ - check(this.dimension == 2) {"Not matrix"} + public operator fun set(i: Int, j: Int, value: T): Unit { + check(this.dimension == 2) { "Not matrix" } this[intArrayOf(i, j)] = value } @@ -31,3 +30,18 @@ public class IntTensor( shape: IntArray, buffer: IntArray ) : BufferedTensor(shape, IntBuffer(buffer)) + +public class LongTensor( + shape: IntArray, + buffer: LongArray +) : BufferedTensor(shape, LongBuffer(buffer)) + +public class FloatTensor( + shape: IntArray, + buffer: FloatArray +) : BufferedTensor(shape, FloatBuffer(buffer)) + +public class RealTensor( + shape: IntArray, + buffer: DoubleArray +) : BufferedTensor(shape, RealBuffer(buffer)) \ No newline at end of file diff --git a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/RealTensorAlgebra.kt b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/RealTensorAlgebra.kt index f91f44519..b3e54f077 100644 --- a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/RealTensorAlgebra.kt +++ b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/RealTensorAlgebra.kt @@ -6,11 +6,6 @@ import kotlin.math.abs import kotlin.math.max -public class RealTensor( - shape: IntArray, - buffer: DoubleArray -) : BufferedTensor(shape, RealBuffer(buffer)) - public class RealTensorAlgebra : TensorPartialDivisionAlgebra { override fun RealTensor.value(): Double { @@ -54,64 +49,6 @@ public class RealTensorAlgebra : TensorPartialDivisionAlgebra { - val totalShape = broadcastShapes(*(tensors.map { it.shape }).toTypedArray()) - val n = totalShape.reduce{ acc, i -> acc * i } - - val res = ArrayList(0) - for (tensor in tensors) { - val resTensor = RealTensor(totalShape, DoubleArray(n)) - - for (linearIndex in 0 until n) { - val totalMultiIndex = resTensor.strides.index(linearIndex) - val curMultiIndex = tensor.shape.copyOf() - - val offset = totalMultiIndex.size - curMultiIndex.size - - for (i in curMultiIndex.indices) { - if (curMultiIndex[i] != 1) { - curMultiIndex[i] = totalMultiIndex[i + offset] - } else { - curMultiIndex[i] = 0 - } - } - - val curLinearIndex = tensor.strides.offset(curMultiIndex) - resTensor.buffer.array[linearIndex] = tensor.buffer.array[curLinearIndex] - } - res.add(resTensor) - } - - return res - } override fun Double.plus(other: RealTensor): RealTensor { val resBuffer = DoubleArray(other.buffer.size) { i -> diff --git a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/TensorAlgebra.kt b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/TensorAlgebra.kt index ec257e45c..280acf8ea 100644 --- a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/TensorAlgebra.kt +++ b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/TensorAlgebra.kt @@ -13,9 +13,6 @@ public interface TensorAlgebra> { public fun TensorType.copy(): TensorType - public fun broadcastShapes(vararg shapes: IntArray): IntArray - public fun broadcastTensors(vararg tensors: RealTensor): List - public operator fun T.plus(other: TensorType): TensorType public operator fun TensorType.plus(value: T): TensorType public operator fun TensorType.plus(other: TensorType): TensorType @@ -87,45 +84,4 @@ public interface TensorPartialDivisionAlgebra //https://pytorch.org/docs/stable/generated/torch.symeig.html public fun TensorType.symEig(eigenvectors: Boolean = true): Pair -} - -public inline fun , - TorchTensorAlgebraType : TensorAlgebra> - TorchTensorAlgebraType.checkShapeCompatible( - a: TensorType, b: TensorType -): Unit = - check(a.shape contentEquals b.shape) { - "Tensors must be of identical shape" - } - -public inline fun , - TorchTensorAlgebraType : TensorAlgebra> - TorchTensorAlgebraType.checkDot(a: TensorType, b: TensorType): Unit { - val sa = a.shape - val sb = b.shape - val na = sa.size - val nb = sb.size - var status: Boolean - if (nb == 1) { - status = sa.last() == sb[0] - } else { - status = sa.last() == sb[nb - 2] - if ((na > 2) and (nb > 2)) { - status = status and - (sa.take(nb - 2).toIntArray() contentEquals sb.take(nb - 2).toIntArray()) - } - } - check(status) { "Incompatible shapes $sa and $sb for dot product" } -} - -public inline fun , - TorchTensorAlgebraType : TensorAlgebra> - TorchTensorAlgebraType.checkTranspose(dim: Int, i: Int, j: Int): Unit = - check((i < dim) and (j < dim)) { - "Cannot transpose $i to $j for a tensor of dim $dim" - } - -public inline fun , - TorchTensorAlgebraType : TensorAlgebra> - TorchTensorAlgebraType.checkView(a: TensorType, shape: IntArray): Unit = - check(a.shape.reduce(Int::times) == shape.reduce(Int::times)) \ No newline at end of file +} \ No newline at end of file diff --git a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/TensorStrides.kt b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/TensorStrides.kt index d6a6f5f16..d1d1204b4 100644 --- a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/TensorStrides.kt +++ b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/TensorStrides.kt @@ -5,7 +5,7 @@ import space.kscience.kmath.nd.offsetFromIndex import kotlin.math.max -public inline fun stridesFromShape(shape: IntArray): IntArray { +internal inline fun stridesFromShape(shape: IntArray): IntArray { val nDim = shape.size val res = IntArray(nDim) if (nDim == 0) @@ -22,7 +22,7 @@ public inline fun stridesFromShape(shape: IntArray): IntArray { } -public inline fun indexFromOffset(offset: Int, strides: IntArray, nDim: Int): IntArray { +internal inline fun indexFromOffset(offset: Int, strides: IntArray, nDim: Int): IntArray { val res = IntArray(nDim) var current = offset var strideIndex = 0 @@ -35,7 +35,7 @@ public inline fun indexFromOffset(offset: Int, strides: IntArray, nDim: Int): In return res } -public inline fun nextIndex(index: IntArray, shape: IntArray, nDim: Int): IntArray { +internal inline fun nextIndex(index: IntArray, shape: IntArray, nDim: Int): IntArray { val res = index.copyOf() var current = nDim - 1 var carry = 0 diff --git a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/utils.kt b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/utils.kt new file mode 100644 index 000000000..65409bb15 --- /dev/null +++ b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/utils.kt @@ -0,0 +1,96 @@ +package space.kscience.kmath.tensors + +import space.kscience.kmath.structures.array +import kotlin.math.max + + +internal inline fun broadcastShapes(vararg shapes: IntArray): IntArray { + var totalDim = 0 + for (shape in shapes) { + totalDim = max(totalDim, shape.size) + } + + val totalShape = IntArray(totalDim) { 0 } + for (shape in shapes) { + for (i in shape.indices) { + val curDim = shape[i] + val offset = totalDim - shape.size + totalShape[i + offset] = max(totalShape[i + offset], curDim) + } + } + + for (shape in shapes) { + for (i in shape.indices) { + val curDim = shape[i] + val offset = totalDim - shape.size + if (curDim != 1 && totalShape[i + offset] != curDim) { + throw RuntimeException("Shapes are not compatible and cannot be broadcast") + } + } + } + + return totalShape +} + +internal inline fun broadcastTensors(vararg tensors: RealTensor): List { + val totalShape = broadcastShapes(*(tensors.map { it.shape }).toTypedArray()) + val n = totalShape.reduce { acc, i -> acc * i } + + val res = ArrayList(0) + for (tensor in tensors) { + val resTensor = RealTensor(totalShape, DoubleArray(n)) + + for (linearIndex in 0 until n) { + val totalMultiIndex = resTensor.strides.index(linearIndex) + val curMultiIndex = tensor.shape.copyOf() + + val offset = totalMultiIndex.size - curMultiIndex.size + + for (i in curMultiIndex.indices) { + if (curMultiIndex[i] != 1) { + curMultiIndex[i] = totalMultiIndex[i + offset] + } else { + curMultiIndex[i] = 0 + } + } + + val curLinearIndex = tensor.strides.offset(curMultiIndex) + resTensor.buffer.array[linearIndex] = tensor.buffer.array[curLinearIndex] + } + res.add(resTensor) + } + + return res +} + +internal inline fun , + TorchTensorAlgebraType : TensorAlgebra> + TorchTensorAlgebraType.checkDot(a: TensorType, b: TensorType): Unit { + val sa = a.shape + val sb = b.shape + val na = sa.size + val nb = sb.size + var status: Boolean + if (nb == 1) { + status = sa.last() == sb[0] + } else { + status = sa.last() == sb[nb - 2] + if ((na > 2) and (nb > 2)) { + status = status and + (sa.take(nb - 2).toIntArray() contentEquals sb.take(nb - 2).toIntArray()) + } + } + check(status) { "Incompatible shapes $sa and $sb for dot product" } +} + +internal inline fun , + TorchTensorAlgebraType : TensorAlgebra> + TorchTensorAlgebraType.checkTranspose(dim: Int, i: Int, j: Int): Unit = + check((i < dim) and (j < dim)) { + "Cannot transpose $i to $j for a tensor of dim $dim" + } + +internal inline fun , + TorchTensorAlgebraType : TensorAlgebra> + TorchTensorAlgebraType.checkView(a: TensorType, shape: IntArray): Unit = + check(a.shape.reduce(Int::times) == shape.reduce(Int::times)) \ No newline at end of file diff --git a/kmath-core/src/commonTest/kotlin/space/kscience/kmath/tensors/TestRealTensorAlgebra.kt b/kmath-core/src/commonTest/kotlin/space/kscience/kmath/tensors/TestRealTensorAlgebra.kt index c9fd6bb3a..7c27f5fa9 100644 --- a/kmath-core/src/commonTest/kotlin/space/kscience/kmath/tensors/TestRealTensorAlgebra.kt +++ b/kmath-core/src/commonTest/kotlin/space/kscience/kmath/tensors/TestRealTensorAlgebra.kt @@ -2,8 +2,6 @@ package space.kscience.kmath.tensors import space.kscience.kmath.structures.array import kotlin.test.Test -import kotlin.test.assertFails -import kotlin.test.assertFailsWith import kotlin.test.assertTrue class TestRealTensorAlgebra { @@ -51,11 +49,11 @@ class TestRealTensorAlgebra { @Test fun broadcastShapes() = RealTensorAlgebra { - assertTrue(this.broadcastShapes( + assertTrue(broadcastShapes( intArrayOf(2, 3), intArrayOf(1, 3), intArrayOf(1, 1, 1) ) contentEquals intArrayOf(1, 2, 3)) - assertTrue(this.broadcastShapes( + assertTrue(broadcastShapes( intArrayOf(6, 7), intArrayOf(5, 6, 1), intArrayOf(7,), intArrayOf(5, 1, 7) ) contentEquals intArrayOf(5, 6, 7)) } @@ -66,7 +64,7 @@ class TestRealTensorAlgebra { val tensor2 = RealTensor(intArrayOf(1, 3), doubleArrayOf(10.0, 20.0, 30.0)) val tensor3 = RealTensor(intArrayOf(1, 1, 1), doubleArrayOf(500.0)) - val res = this.broadcastTensors(tensor1, tensor2, tensor3) + val res = broadcastTensors(tensor1, tensor2, tensor3) assertTrue(res[0].shape contentEquals intArrayOf(1, 2, 3)) assertTrue(res[1].shape contentEquals intArrayOf(1, 2, 3))