KMP library for tensors #300

Merged
grinisrit merged 215 commits from feature/tensor-algebra into dev 2021-05-08 09:48:04 +03:00
10 changed files with 361 additions and 24 deletions
Showing only changes of commit 7c38b16f17 - Show all commits

5
.gitignore vendored
View File

@ -10,5 +10,6 @@ out/
# Cache of project
.gradletasknamecache
# Generated by javac -h
*.class
# Generated by javac -h and runtime
*.class
*.log

View File

@ -162,7 +162,7 @@ public interface Strides {
/**
* Array strides
*/
public val strides: List<Int>
public val strides: IntArray
/**
* Get linear index from multidimensional index
@ -189,6 +189,11 @@ 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].
*/
@ -199,7 +204,7 @@ public class DefaultStrides private constructor(override val shape: IntArray) :
/**
* Strides for memory access
*/
override val strides: List<Int> by lazy {
override val strides: IntArray by lazy {
sequence {
var current = 1
yield(1)
@ -208,13 +213,10 @@ public class DefaultStrides private constructor(override val shape: IntArray) :
current *= it
yield(current)
}
}.toList()
}.toList().toIntArray()
}
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,${this.shape[i]})")
value * strides[i]
}.sum()
override fun offset(index: IntArray): Int = offsetFromIndex(index, shape, strides)
override fun index(offset: Int): IntArray {
val res = IntArray(shape.size)
@ -322,7 +324,7 @@ public inline fun <T, reified R : Any> NDStructure<T>.mapToBuffer(
/**
* Mutable ND buffer based on linear [MutableBuffer].
*/
public class MutableNDBuffer<T>(
public open class MutableNDBuffer<T>(
strides: Strides,
buffer: MutableBuffer<T>,
) : NDBuffer<T>(strides, buffer), MutableNDStructure<T> {

View File

@ -175,6 +175,7 @@ public interface SpaceOperations<T> : Algebra<T> {
* @param k the divisor.
* @return the quotient.
*/
@Deprecated("Dividing not allowed in a Ring")
public operator fun T.div(k: Number): T = multiply(this, 1.0 / k.toDouble())
/**

View File

@ -0,0 +1,173 @@
package space.kscience.kmath.tensors
import space.kscience.kmath.nd.MutableNDBuffer
import space.kscience.kmath.structures.RealBuffer
import space.kscience.kmath.structures.array
public class RealTensor(
override val shape: IntArray,
buffer: DoubleArray
) :
TensorStructure<Double>,
MutableNDBuffer<Double>(
TensorStrides(shape),
RealBuffer(buffer)
) {
override fun item(): Double = buffer[0]
}
public class RealTensorAlgebra : TensorPartialDivisionAlgebra<Double, RealTensor> {
override fun add(a: RealTensor, b: RealTensor): RealTensor {
TODO("Not yet implemented")
}
override fun multiply(a: RealTensor, k: Number): RealTensor {
TODO("Not yet implemented")
}
override val zero: RealTensor
get() = TODO("Not yet implemented")
override fun multiply(a: RealTensor, b: RealTensor): RealTensor {
TODO("Not yet implemented")
}
override val one: RealTensor
get() = TODO("Not yet implemented")
override fun Double.plus(other: RealTensor): RealTensor {
val n = other.buffer.size
val arr = other.buffer.array
val res = DoubleArray(n)
for (i in 1..n)
res[i - 1] = arr[i - 1] + this
return RealTensor(other.shape, res)
}
override fun RealTensor.plus(value: Double): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.plusAssign(value: Double) {
TODO("Not yet implemented")
}
override fun RealTensor.plusAssign(other: RealTensor) {
TODO("Not yet implemented")
}
override fun Double.minus(other: RealTensor): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.minus(value: Double): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.minusAssign(value: Double) {
TODO("Not yet implemented")
}
override fun RealTensor.minusAssign(other: RealTensor) {
TODO("Not yet implemented")
}
override fun Double.times(other: RealTensor): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.times(value: Double): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.timesAssign(value: Double) {
TODO("Not yet implemented")
}
override fun RealTensor.timesAssign(other: RealTensor) {
TODO("Not yet implemented")
}
override fun RealTensor.dot(other: RealTensor): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.dotAssign(other: RealTensor) {
TODO("Not yet implemented")
}
override fun RealTensor.dotRightAssign(other: RealTensor) {
TODO("Not yet implemented")
}
override fun diagonalEmbedding(diagonalEntries: RealTensor, offset: Int, dim1: Int, dim2: Int): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.transpose(i: Int, j: Int): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.transposeAssign(i: Int, j: Int) {
TODO("Not yet implemented")
}
override fun RealTensor.view(shape: IntArray): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.abs(): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.absAssign() {
TODO("Not yet implemented")
}
override fun RealTensor.sum(): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.sumAssign() {
TODO("Not yet implemented")
}
override fun RealTensor.div(other: RealTensor): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.divAssign(other: RealTensor) {
TODO("Not yet implemented")
}
override fun RealTensor.exp(): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.expAssign() {
TODO("Not yet implemented")
}
override fun RealTensor.log(): RealTensor {
TODO("Not yet implemented")
}
override fun RealTensor.logAssign() {
TODO("Not yet implemented")
}
override fun RealTensor.svd(): Triple<RealTensor, RealTensor, RealTensor> {
TODO("Not yet implemented")
}
override fun RealTensor.symEig(eigenvectors: Boolean): Pair<RealTensor, RealTensor> {
TODO("Not yet implemented")
}
}
public inline fun <R> RealTensorAlgebra(block: RealTensorAlgebra.() -> R): R =
RealTensorAlgebra().block()

View File

@ -1,70 +1,114 @@
package space.kscience.kmath.tensors
import space.kscience.kmath.nd.MutableNDStructure
public interface TensorStructure<T> : MutableNDStructure<T> {
// A tensor can have empty shape, in which case it represents just a value
public fun value(): T
}
import space.kscience.kmath.operations.Ring
import space.kscience.kmath.operations.RingWithNumbers
// https://proofwiki.org/wiki/Definition:Algebra_over_Ring
public interface TensorAlgebra<T, TensorType : TensorStructure<T>> {
public interface TensorAlgebra<T, TensorType : TensorStructure<T>>: RingWithNumbers<TensorType> {
public operator fun T.plus(other: TensorType): TensorType
public operator fun TensorType.plus(value: T): TensorType
public operator fun TensorType.plus(other: TensorType): TensorType
public operator fun TensorType.plusAssign(value: T): Unit
public operator fun TensorType.plusAssign(other: TensorType): Unit
public operator fun T.minus(other: TensorType): TensorType
public operator fun TensorType.minus(value: T): TensorType
public operator fun TensorType.minus(other: TensorType): TensorType
public operator fun TensorType.minusAssign(value: T): Unit
public operator fun TensorType.minusAssign(other: TensorType): Unit
public operator fun T.times(other: TensorType): TensorType
public operator fun TensorType.times(value: T): TensorType
public operator fun TensorType.times(other: TensorType): TensorType
public operator fun TensorType.timesAssign(value: T): Unit
public operator fun TensorType.timesAssign(other: TensorType): Unit
public operator fun TensorType.unaryMinus(): TensorType
//https://pytorch.org/docs/stable/generated/torch.matmul.html
public infix fun TensorType.dot(other: TensorType): TensorType
public infix fun TensorType.dotAssign(other: TensorType): Unit
public infix fun TensorType.dotRightAssign(other: TensorType): Unit
//https://pytorch.org/docs/stable/generated/torch.diag_embed.html
public fun diagonalEmbedding(
diagonalEntries: TensorType,
offset: Int = 0, dim1: Int = -2, dim2: Int = -1
): TensorType
//https://pytorch.org/docs/stable/generated/torch.transpose.html
public fun TensorType.transpose(i: Int, j: Int): TensorType
public fun TensorType.transposeAssign(i: Int, j: Int): Unit
//https://pytorch.org/docs/stable/tensor_view.html
public fun TensorType.view(shape: IntArray): TensorType
//https://pytorch.org/docs/stable/generated/torch.abs.html
public fun TensorType.abs(): TensorType
public fun TensorType.absAssign(): Unit
//https://pytorch.org/docs/stable/generated/torch.sum.html
public fun TensorType.sum(): TensorType
public fun TensorType.sumAssign(): Unit
}
// https://proofwiki.org/wiki/Definition:Division_Algebra
public interface TensorPartialDivisionAlgebra<T, TensorType : TensorStructure<T>> :
TensorAlgebra<T, TensorType> {
public operator fun TensorType.div(other: TensorType): TensorType
public operator fun TensorType.divAssign(other: TensorType)
//https://pytorch.org/docs/stable/generated/torch.exp.html
public fun TensorType.exp(): TensorType
public fun TensorType.expAssign(): Unit
//https://pytorch.org/docs/stable/generated/torch.log.html
public fun TensorType.log(): TensorType
public fun TensorType.logAssign(): Unit
//https://pytorch.org/docs/stable/generated/torch.svd.html
public fun TensorType.svd(): Triple<TensorType, TensorType, TensorType>
//https://pytorch.org/docs/stable/generated/torch.symeig.html
public fun TensorType.symEig(eigenvectors: Boolean = true): Pair<TensorType, TensorType>
}
}
public inline fun <T, TensorType : TensorStructure<T>,
TorchTensorAlgebraType : TensorAlgebra<T, TensorType>>
TorchTensorAlgebraType.checkShapeCompatible(
a: TensorType, b: TensorType
): Unit =
check(a.shape contentEquals b.shape) {
"Tensors must be of identical shape"
}
public inline fun <T, TensorType : TensorStructure<T>,
TorchTensorAlgebraType : TensorAlgebra<T, TensorType>>
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 <T, TensorType : TensorStructure<T>,
TorchTensorAlgebraType : TensorAlgebra<T, TensorType>>
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 <T, TensorType : TensorStructure<T>,
TorchTensorAlgebraType : TensorAlgebra<T, TensorType>>
TorchTensorAlgebraType.checkView(a: TensorType, shape: IntArray): Unit =
check(a.shape.reduce(Int::times) == shape.reduce(Int::times))

View File

@ -0,0 +1,52 @@
package space.kscience.kmath.tensors
import space.kscience.kmath.nd.Strides
import space.kscience.kmath.nd.offsetFromIndex
import kotlin.math.max
inline public fun stridesFromShape(shape: IntArray): IntArray {
val nDim = shape.size
val res = IntArray(nDim)
if (nDim == 0)
return res
var current = nDim - 1
res[current] = 1
while (current > 0) {
res[current - 1] = max(1, shape[current]) * res[current]
current--
}
return res
}
inline public fun indexFromOffset(offset: Int, strides: IntArray, nDim: Int): IntArray {
val res = IntArray(nDim)
var current = offset
var strideIndex = 0
while (strideIndex < nDim) {
res[strideIndex] = (current / strides[strideIndex])
current %= strides[strideIndex]
strideIndex++
}
return res
}
public class TensorStrides(override val shape: IntArray): Strides
{
override val strides: IntArray
get() = stridesFromShape(shape)
override fun offset(index: IntArray): Int = offsetFromIndex(index, shape, strides)
override fun index(offset: Int): IntArray =
indexFromOffset(offset, strides, shape.size)
override val linearSize: Int
get() = shape.fold(1) { acc, i -> acc * i }
}

View File

@ -0,0 +1,23 @@
package space.kscience.kmath.tensors
import space.kscience.kmath.nd.MutableNDStructure
public interface TensorStructure<T> : MutableNDStructure<T> {
public fun item(): T
// A tensor can have empty shape, in which case it represents just a value
public fun value(): T {
checkIsValue()
return item()
}
}
public inline fun <T> TensorStructure<T>.isValue(): Boolean {
return (dimension == 0)
}
public inline fun <T> TensorStructure<T>.isNotValue(): Boolean = !this.isValue()
public inline fun <T> TensorStructure<T>.checkIsValue(): Unit = check(this.isValue()) {
"This tensor has shape ${shape.toList()}"
}

View File

@ -0,0 +1,24 @@
package space.kscience.kmath.tensors
import space.kscience.kmath.structures.array
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertTrue
class TestRealTensor {
@Test
fun valueTest(){
val value = 12.5
val tensor = RealTensor(IntArray(0), doubleArrayOf(value))
assertEquals(tensor.value(), value)
}
@Test
fun stridesTest(){
val tensor = RealTensor(intArrayOf(2,2), doubleArrayOf(3.5,5.8,58.4,2.4))
assertEquals(tensor[intArrayOf(0,1)], 5.8)
assertTrue(tensor.elements().map{ it.second }.toList().toDoubleArray() contentEquals tensor.buffer.array)
}
}

View File

@ -0,0 +1,16 @@
package space.kscience.kmath.tensors
import space.kscience.kmath.structures.array
import kotlin.test.Test
import kotlin.test.assertTrue
class TestRealTensorAlgebra {
@Test
fun doublePlus() = RealTensorAlgebra {
val tensor = RealTensor(intArrayOf(2), doubleArrayOf(1.0, 2.0))
val res = 10.0 + tensor
assertTrue(res.buffer.array contentEquals doubleArrayOf(11.0,12.0))
}
}

View File

@ -96,6 +96,7 @@ public interface Nd4jArraySpace<T, S : Space<T>> : NDSpace<T, S>, Nd4jArrayAlgeb
return a.ndArray.mul(k).wrap()
}
@Deprecated("Avoid using this method, underlying array get casted to Doubles")
public override operator fun NDStructure<T>.div(k: Number): Nd4jArrayStructure<T> {
return ndArray.div(k).wrap()
}