diff --git a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/nd/StructureND.kt b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/nd/StructureND.kt index 86b867130..f28d29e47 100644 --- a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/nd/StructureND.kt +++ b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/nd/StructureND.kt @@ -177,7 +177,7 @@ public interface Strides { /** * Array strides */ - public val strides: IntArray + public val strides: List /** * Get linear index from multidimensional index @@ -209,12 +209,6 @@ public interface Strides { } } -internal inline fun offsetFromIndex(index: IntArray, shape: IntArray, strides: IntArray): Int = - index.mapIndexed { i, value -> - if (value < 0 || value >= shape[i]) throw IndexOutOfBoundsException("Index $value out of shape bounds: (0,${shape[i]})") - value * strides[i] - }.sum() - /** * Simple implementation of [Strides]. */ @@ -225,7 +219,7 @@ public class DefaultStrides private constructor(override val shape: IntArray) : /** * Strides for memory access */ - override val strides: IntArray by lazy { + override val strides: List by lazy { sequence { var current = 1 yield(1) @@ -234,10 +228,13 @@ public class DefaultStrides private constructor(override val shape: IntArray) : current *= it yield(current) } - }.toList().toIntArray() + }.toList() } - override fun offset(index: IntArray): Int = offsetFromIndex(index, shape, strides) + override fun offset(index: IntArray): Int = index.mapIndexed { i, value -> + if (value < 0 || value >= shape[i]) throw IndexOutOfBoundsException("Index $value out of shape bounds: (0,${shape[i]})") + value * strides[i] + }.sum() override fun index(offset: Int): IntArray { val res = IntArray(shape.size) diff --git a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/core/BroadcastDoubleTensorAlgebra.kt b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/core/BroadcastDoubleTensorAlgebra.kt index 20f64f469..4bfbe5863 100644 --- a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/core/BroadcastDoubleTensorAlgebra.kt +++ b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/core/BroadcastDoubleTensorAlgebra.kt @@ -8,7 +8,7 @@ public class BroadcastDoubleTensorAlgebra : DoubleTensorAlgebra() { val broadcast = broadcastTensors(this, other) val newThis = broadcast[0] val newOther = broadcast[1] - val resBuffer = DoubleArray(newThis.strides.linearSize) { i -> + val resBuffer = DoubleArray(newThis.linearStructure.size) { i -> newThis.buffer.array()[i] + newOther.buffer.array()[i] } return DoubleTensor(newThis.shape, resBuffer) @@ -16,7 +16,7 @@ public class BroadcastDoubleTensorAlgebra : DoubleTensorAlgebra() { override fun DoubleTensor.plusAssign(other: DoubleTensor) { val newOther = broadcastTo(other, this.shape) - for (i in 0 until this.strides.linearSize) { + for (i in 0 until this.linearStructure.size) { this.buffer.array()[this.bufferStart + i] += newOther.buffer.array()[this.bufferStart + i] } @@ -26,7 +26,7 @@ public class BroadcastDoubleTensorAlgebra : DoubleTensorAlgebra() { val broadcast = broadcastTensors(this, other) val newThis = broadcast[0] val newOther = broadcast[1] - val resBuffer = DoubleArray(newThis.strides.linearSize) { i -> + val resBuffer = DoubleArray(newThis.linearStructure.size) { i -> newThis.buffer.array()[i] - newOther.buffer.array()[i] } return DoubleTensor(newThis.shape, resBuffer) @@ -34,7 +34,7 @@ public class BroadcastDoubleTensorAlgebra : DoubleTensorAlgebra() { override fun DoubleTensor.minusAssign(other: DoubleTensor) { val newOther = broadcastTo(other, this.shape) - for (i in 0 until this.strides.linearSize) { + for (i in 0 until this.linearStructure.size) { this.buffer.array()[this.bufferStart + i] -= newOther.buffer.array()[this.bufferStart + i] } @@ -44,7 +44,7 @@ public class BroadcastDoubleTensorAlgebra : DoubleTensorAlgebra() { val broadcast = broadcastTensors(this, other) val newThis = broadcast[0] val newOther = broadcast[1] - val resBuffer = DoubleArray(newThis.strides.linearSize) { i -> + val resBuffer = DoubleArray(newThis.linearStructure.size) { i -> newThis.buffer.array()[newOther.bufferStart + i] * newOther.buffer.array()[newOther.bufferStart + i] } @@ -53,7 +53,7 @@ public class BroadcastDoubleTensorAlgebra : DoubleTensorAlgebra() { override fun DoubleTensor.timesAssign(other: DoubleTensor) { val newOther = broadcastTo(other, this.shape) - for (i in 0 until this.strides.linearSize) { + for (i in 0 until this.linearStructure.size) { this.buffer.array()[this.bufferStart + i] *= newOther.buffer.array()[this.bufferStart + i] } @@ -63,7 +63,7 @@ public class BroadcastDoubleTensorAlgebra : DoubleTensorAlgebra() { val broadcast = broadcastTensors(this, other) val newThis = broadcast[0] val newOther = broadcast[1] - val resBuffer = DoubleArray(newThis.strides.linearSize) { i -> + val resBuffer = DoubleArray(newThis.linearStructure.size) { i -> newThis.buffer.array()[newOther.bufferStart + i] / newOther.buffer.array()[newOther.bufferStart + i] } @@ -72,7 +72,7 @@ public class BroadcastDoubleTensorAlgebra : DoubleTensorAlgebra() { override fun DoubleTensor.divAssign(other: DoubleTensor) { val newOther = broadcastTo(other, this.shape) - for (i in 0 until this.strides.linearSize) { + for (i in 0 until this.linearStructure.size) { this.buffer.array()[this.bufferStart + i] /= newOther.buffer.array()[this.bufferStart + i] } @@ -130,7 +130,7 @@ internal inline fun broadcastTo(tensor: DoubleTensor, newShape: IntArray): Doubl } for (linearIndex in 0 until n) { - val totalMultiIndex = resTensor.strides.index(linearIndex) + val totalMultiIndex = resTensor.linearStructure.index(linearIndex) val curMultiIndex = tensor.shape.copyOf() val offset = totalMultiIndex.size - curMultiIndex.size @@ -143,7 +143,7 @@ internal inline fun broadcastTo(tensor: DoubleTensor, newShape: IntArray): Doubl } } - val curLinearIndex = tensor.strides.offset(curMultiIndex) + val curLinearIndex = tensor.linearStructure.offset(curMultiIndex) resTensor.buffer.array()[linearIndex] = tensor.buffer.array()[tensor.bufferStart + curLinearIndex] } @@ -159,7 +159,7 @@ internal inline fun broadcastTensors(vararg tensors: DoubleTensor): List( internal val bufferStart: Int ) : TensorStructure { - public val strides: TensorStrides - get() = TensorStrides(shape) + public val linearStructure: TensorLinearStructure + get() = TensorLinearStructure(shape) public val numel: Int - get() = strides.linearSize + get() = linearStructure.size - override fun get(index: IntArray): T = buffer[bufferStart + strides.offset(index)] + override fun get(index: IntArray): T = buffer[bufferStart + linearStructure.offset(index)] override fun set(index: IntArray, value: T) { - buffer[bufferStart + strides.offset(index)] = value + buffer[bufferStart + linearStructure.offset(index)] = value } - override fun elements(): Sequence> = strides.indices().map { + override fun elements(): Sequence> = linearStructure.indices().map { it to this[it] } @@ -35,7 +33,7 @@ public open class BufferedTensor( public fun vectorSequence(): Sequence> = sequence { check(shape.size >= 1) {"todo"} - val vectorOffset = strides.strides[0] + val vectorOffset = linearStructure.strides[0] val vectorShape = intArrayOf(shape.last()) for (offset in 0 until numel step vectorOffset) { val vector = BufferedTensor(vectorShape, buffer, offset).as1D() @@ -45,7 +43,7 @@ public open class BufferedTensor( public fun matrixSequence(): Sequence> = sequence { check(shape.size >= 2) {"todo"} - val matrixOffset = strides.strides[1] + val matrixOffset = linearStructure.strides[1] val matrixShape = intArrayOf(shape[shape.size - 2], shape.last()) //todo better way? for (offset in 0 until numel step matrixOffset) { val matrix = BufferedTensor(matrixShape, buffer, offset).as2D() diff --git a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt index 875f21687..385118d97 100644 --- a/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt +++ b/kmath-core/src/commonMain/kotlin/space/kscience/kmath/tensors/core/DoubleTensorAlgebra.kt @@ -27,7 +27,7 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra + val resBuffer = DoubleArray(other.linearStructure.size) { i -> other.buffer.array()[other.bufferStart + i] + this } return DoubleTensor(other.shape, resBuffer) @@ -64,35 +64,35 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra + val resBuffer = DoubleArray(this.linearStructure.size) { i -> this.buffer.array()[i] + other.buffer.array()[i] } return DoubleTensor(this.shape, resBuffer) } override fun DoubleTensor.plusAssign(value: Double) { - for (i in 0 until this.strides.linearSize) { + for (i in 0 until this.linearStructure.size) { this.buffer.array()[this.bufferStart + i] += value } } override fun DoubleTensor.plusAssign(other: DoubleTensor) { checkShapesCompatible(this, other) - for (i in 0 until this.strides.linearSize) { + for (i in 0 until this.linearStructure.size) { this.buffer.array()[this.bufferStart + i] += other.buffer.array()[this.bufferStart + i] } } override fun Double.minus(other: DoubleTensor): DoubleTensor { - val resBuffer = DoubleArray(other.strides.linearSize) { i -> + val resBuffer = DoubleArray(other.linearStructure.size) { i -> this - other.buffer.array()[other.bufferStart + i] } return DoubleTensor(other.shape, resBuffer) } override fun DoubleTensor.minus(value: Double): DoubleTensor { - val resBuffer = DoubleArray(this.strides.linearSize) { i -> + val resBuffer = DoubleArray(this.linearStructure.size) { i -> this.buffer.array()[this.bufferStart + i] - value } return DoubleTensor(this.shape, resBuffer) @@ -100,28 +100,28 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra + val resBuffer = DoubleArray(this.linearStructure.size) { i -> this.buffer.array()[i] - other.buffer.array()[i] } return DoubleTensor(this.shape, resBuffer) } override fun DoubleTensor.minusAssign(value: Double) { - for (i in 0 until this.strides.linearSize) { + for (i in 0 until this.linearStructure.size) { this.buffer.array()[this.bufferStart + i] -= value } } override fun DoubleTensor.minusAssign(other: DoubleTensor) { checkShapesCompatible(this, other) - for (i in 0 until this.strides.linearSize) { + for (i in 0 until this.linearStructure.size) { this.buffer.array()[this.bufferStart + i] -= other.buffer.array()[this.bufferStart + i] } } override fun Double.times(other: DoubleTensor): DoubleTensor { - val resBuffer = DoubleArray(other.strides.linearSize) { i -> + val resBuffer = DoubleArray(other.linearStructure.size) { i -> other.buffer.array()[other.bufferStart + i] * this } return DoubleTensor(other.shape, resBuffer) @@ -131,7 +131,7 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra + val resBuffer = DoubleArray(this.linearStructure.size) { i -> this.buffer.array()[other.bufferStart + i] * other.buffer.array()[other.bufferStart + i] } @@ -139,21 +139,21 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra + val resBuffer = DoubleArray(this.linearStructure.size) { i -> this.buffer.array()[this.bufferStart + i] / value } return DoubleTensor(this.shape, resBuffer) @@ -161,7 +161,7 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra + val resBuffer = DoubleArray(this.linearStructure.size) { i -> this.buffer.array()[other.bufferStart + i] / other.buffer.array()[other.bufferStart + i] } @@ -169,21 +169,21 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra + val resBuffer = DoubleArray(this.linearStructure.size) { i -> this.buffer.array()[this.bufferStart + i].unaryMinus() } return DoubleTensor(this.shape, resBuffer) @@ -191,7 +191,7 @@ public open class DoubleTensorAlgebra : TensorPartialDivisionAlgebra + if (value < 0 || value >= shape[i]) throw IndexOutOfBoundsException("Index $value out of shape bounds: (0,${shape[i]})") + value * strides[i] + }.sum() + internal inline fun stridesFromShape(shape: IntArray): IntArray { val nDim = shape.size val res = IntArray(nDim) @@ -35,7 +39,7 @@ internal inline fun indexFromOffset(offset: Int, strides: IntArray, nDim: Int): return res } -internal inline fun nextIndex(index: IntArray, shape: IntArray, nDim: Int): IntArray { +internal inline fun stepIndex(index: IntArray, shape: IntArray, nDim: Int): IntArray { val res = index.copyOf() var current = nDim - 1 var carry = 0 @@ -47,26 +51,29 @@ internal inline fun nextIndex(index: IntArray, shape: IntArray, nDim: Int): IntA res[current] = 0 } current-- - } while(carry != 0 && current >= 0) + } while (carry != 0 && current >= 0) return res } - -public class TensorStrides(override val shape: IntArray): Strides +public class TensorLinearStructure(public val shape: IntArray) { - override val strides: IntArray + public val strides: IntArray get() = stridesFromShape(shape) - override fun offset(index: IntArray): Int = offsetFromIndex(index, shape, strides) + public fun offset(index: IntArray): Int = offsetFromIndex(index, shape, strides) - override fun index(offset: Int): IntArray = + public fun index(offset: Int): IntArray = indexFromOffset(offset, strides, shape.size) - override fun nextIndex(index: IntArray): IntArray = - nextIndex(index, shape, shape.size) + public fun stepIndex(index: IntArray): IntArray = + stepIndex(index, shape, shape.size) - override val linearSize: Int + public val size: Int get() = shape.reduce(Int::times) + + public fun indices(): Sequence = (0 until size).asSequence().map { + index(it) + } } \ No newline at end of file