diff --git a/kmath-tensorflow/build.gradle.kts b/kmath-tensorflow/build.gradle.kts index fa77a272a..c8307f01f 100644 --- a/kmath-tensorflow/build.gradle.kts +++ b/kmath-tensorflow/build.gradle.kts @@ -7,6 +7,7 @@ description = "Google tensorflow connector" dependencies { api(project(":kmath-tensors")) api("org.tensorflow:tensorflow-core-api:0.3.3") + testImplementation("org.tensorflow:tensorflow-core-platform:0.3.3") } readme { diff --git a/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/DoubleTensorFlowAlgebra.kt b/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/DoubleTensorFlowAlgebra.kt index 44420cd76..eb8245944 100644 --- a/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/DoubleTensorFlowAlgebra.kt +++ b/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/DoubleTensorFlowAlgebra.kt @@ -5,6 +5,7 @@ import org.tensorflow.Output import org.tensorflow.ndarray.NdArray import org.tensorflow.op.core.Constant import org.tensorflow.types.TFloat64 +import space.kscience.kmath.expressions.Symbol import space.kscience.kmath.misc.PerformancePitfall import space.kscience.kmath.nd.DefaultStrides import space.kscience.kmath.nd.Shape @@ -13,20 +14,22 @@ import space.kscience.kmath.operations.DoubleField public class DoubleTensorFlowOutput( graph: Graph, - output: Output + output: Output, ) : TensorFlowOutput(graph, output) { - override fun org.tensorflow.Tensor.actualizeTensor(): NdArray = output.asTensor() + + override fun org.tensorflow.Tensor.actualizeTensor(): NdArray = this as TFloat64 + } public class DoubleTensorFlowAlgebra internal constructor( - graph: Graph + graph: Graph, ) : TensorFlowAlgebra(graph) { override val elementAlgebra: DoubleField get() = DoubleField override fun structureND( shape: Shape, - initializer: DoubleField.(IntArray) -> Double + initializer: DoubleField.(IntArray) -> Double, ): StructureND { val res = TFloat64.tensorOf(org.tensorflow.ndarray.Shape.of(*shape.toLongArray())) { array -> DefaultStrides(shape).forEach { index -> @@ -53,4 +56,20 @@ public class DoubleTensorFlowAlgebra internal constructor( override fun Output.wrap(): TensorFlowOutput = DoubleTensorFlowOutput(graph, this) override fun const(value: Double): Constant = ops.constant(value) + + +} + +public fun DoubleField.produceWithTF( + block: DoubleTensorFlowAlgebra.() -> StructureND, +): StructureND = Graph().use { graph -> + val scope = DoubleTensorFlowAlgebra(graph) + scope.export(scope.block()) +} + +public fun DoubleField.produceMapWithTF( + block: DoubleTensorFlowAlgebra.() -> Map>, +): Map> = Graph().use { graph -> + val scope = DoubleTensorFlowAlgebra(graph) + scope.block().mapValues { scope.export(it.value) } } \ No newline at end of file diff --git a/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/IntTensorFlowAlgebra.kt b/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/IntTensorFlowAlgebra.kt new file mode 100644 index 000000000..084a445e0 --- /dev/null +++ b/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/IntTensorFlowAlgebra.kt @@ -0,0 +1,21 @@ +package space.kscience.kmath.tensorflow + +import org.tensorflow.Graph +import org.tensorflow.Output +import org.tensorflow.ndarray.NdArray +import org.tensorflow.types.TInt32 +import org.tensorflow.types.TInt64 + +public class IntTensorFlowOutput( + graph: Graph, + output: Output, +) : TensorFlowOutput(graph, output) { + override fun org.tensorflow.Tensor.actualizeTensor(): NdArray = this as TInt32 +} + +public class LongTensorFlowOutput( + graph: Graph, + output: Output, +) : TensorFlowOutput(graph, output) { + override fun org.tensorflow.Tensor.actualizeTensor(): NdArray = this as TInt64 +} \ No newline at end of file diff --git a/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/TensorFlowAlgebra.kt b/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/TensorFlowAlgebra.kt index e73620d01..1f978c0fc 100644 --- a/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/TensorFlowAlgebra.kt +++ b/kmath-tensorflow/src/main/kotlin/space/kscience/kmath/tensorflow/TensorFlowAlgebra.kt @@ -8,8 +8,13 @@ import org.tensorflow.Session import org.tensorflow.ndarray.NdArray import org.tensorflow.op.Ops import org.tensorflow.op.core.Constant +import org.tensorflow.op.core.Max +import org.tensorflow.op.core.Min +import org.tensorflow.op.core.Sum +import org.tensorflow.types.TInt32 import org.tensorflow.types.family.TType import space.kscience.kmath.misc.PerformancePitfall +import space.kscience.kmath.misc.UnstableKMathAPI import space.kscience.kmath.nd.Shape import space.kscience.kmath.nd.StructureND import space.kscience.kmath.operations.Ring @@ -38,8 +43,8 @@ public value class TensorFlowArray(public val tensor: NdArray) : Tensor } public abstract class TensorFlowOutput( - private val graph: Graph, - output: Output + protected val graph: Graph, + output: Output, ) : TensorFlowTensor { public var output: Output = output @@ -49,9 +54,10 @@ public abstract class TensorFlowOutput( protected abstract fun org.tensorflow.Tensor.actualizeTensor(): NdArray - private val actualTensor by lazy { - val session = Session(graph) - TensorFlowArray(session.runner().fetch(output).run().first().actualizeTensor()) + internal val actualTensor by lazy { + Session(graph).use { session -> + TensorFlowArray(session.runner().fetch(output).run().first().actualizeTensor()) + } } override fun get(index: IntArray): T = actualTensor[index] @@ -66,9 +72,9 @@ public abstract class TensorFlowOutput( } -public abstract class TensorFlowAlgebra> internal constructor( - protected val graph: Graph -) : TensorAlgebra { +public abstract class TensorFlowAlgebra> internal constructor( + protected val graph: Graph, +) : TensorAlgebra { protected val ops: Ops by lazy { Ops.create(graph) } @@ -83,7 +89,7 @@ public abstract class TensorFlowAlgebra> internal cons private inline fun StructureND.biOp( other: StructureND, - operation: (left: Operand, right: Operand) -> Operand + operation: (left: Operand, right: Operand) -> Operand, ): TensorFlowOutput { val left = asTensorFlow().output val right = other.asTensorFlow().output @@ -92,7 +98,7 @@ public abstract class TensorFlowAlgebra> internal cons private inline fun T.biOp( other: StructureND, - operation: (left: Operand, right: Operand) -> Operand + operation: (left: Operand, right: Operand) -> Operand, ): TensorFlowOutput { val left = const(this) val right = other.asTensorFlow().output @@ -101,7 +107,7 @@ public abstract class TensorFlowAlgebra> internal cons private inline fun StructureND.biOp( value: T, - operation: (left: Operand, right: Operand) -> Operand + operation: (left: Operand, right: Operand) -> Operand, ): TensorFlowOutput { val left = asTensorFlow().output val right = const(value) @@ -110,7 +116,7 @@ public abstract class TensorFlowAlgebra> internal cons private inline fun Tensor.inPlaceOp( other: StructureND, - operation: (left: Operand, right: Operand) -> Operand + operation: (left: Operand, right: Operand) -> Operand, ): Unit { val origin = asTensorFlow() val left = origin.output @@ -120,7 +126,7 @@ public abstract class TensorFlowAlgebra> internal cons private inline fun Tensor.inPlaceOp( value: T, - operation: (left: Operand, right: Operand) -> Operand + operation: (left: Operand, right: Operand) -> Operand, ): Unit { val origin = asTensorFlow() val left = origin.output @@ -128,8 +134,8 @@ public abstract class TensorFlowAlgebra> internal cons origin.output = operation(left, right).asOutput() } - private inline fun unOp(value: StructureND, operation: (Operand) -> Operand): TensorFlowOutput = - operation(value.asTensorFlow().output).asOutput().wrap() + private inline fun StructureND.unOp(operation: (Operand) -> Operand): TensorFlowOutput = + operation(asTensorFlow().output).asOutput().wrap() override fun T.plus(arg: StructureND): TensorFlowOutput = biOp(arg, ops.math::add) @@ -149,67 +155,79 @@ public abstract class TensorFlowAlgebra> internal cons override fun Tensor.minusAssign(value: T): Unit = inPlaceOp(value, ops.math::sub) - override fun Tensor.minusAssign(other: StructureND): Unit = inPlaceOp(other, ops.math::sub) + override fun Tensor.minusAssign(arg: StructureND): Unit = inPlaceOp(arg, ops.math::sub) override fun T.times(arg: StructureND): TensorFlowOutput = biOp(arg, ops.math::mul) override fun StructureND.times(arg: T): TensorFlowOutput = biOp(arg, ops.math::mul) - override fun StructureND.times(other: StructureND): TensorFlowOutput = biOp(other, ops.math::mul) + override fun StructureND.times(arg: StructureND): TensorFlowOutput = biOp(arg, ops.math::mul) override fun Tensor.timesAssign(value: T): Unit = inPlaceOp(value, ops.math::mul) override fun Tensor.timesAssign(arg: StructureND): Unit = inPlaceOp(arg, ops.math::mul) - override fun StructureND.unaryMinus(): TensorFlowOutput = unOp(this, ops.math::neg) + override fun StructureND.unaryMinus(): TensorFlowOutput = unOp(ops.math::neg) - override fun StructureND.get(i: Int): Tensor { + override fun Tensor.get(i: Int): Tensor = unOp { TODO("Not yet implemented") } - override fun StructureND.transpose(i: Int, j: Int): Tensor = unOp(this) { + override fun Tensor.transpose(i: Int, j: Int): Tensor = unOp { ops.linalg.transpose(it, ops.constant(intArrayOf(i, j))) } - override fun Tensor.view(shape: IntArray): Tensor { - TODO("Not yet implemented") + override fun Tensor.view(shape: IntArray): Tensor = unOp { + ops.reshape(it, ops.constant(shape)) } - override fun Tensor.viewAs(other: StructureND): Tensor { - TODO("Not yet implemented") + override fun Tensor.viewAs(other: StructureND): Tensor = biOp(other) { l, r -> + ops.reshape(l, ops.shape(r)) } override fun StructureND.dot(other: StructureND): TensorFlowOutput = biOp(other) { l, r -> ops.linalg.matMul(l, r) } - override fun diagonalEmbedding(diagonalEntries: Tensor, offset: Int, dim1: Int, dim2: Int): Tensor = ops.run { - TODO("Not yet implemented") + override fun diagonalEmbedding( + diagonalEntries: Tensor, + offset: Int, + dim1: Int, + dim2: Int, + ): TensorFlowOutput = diagonalEntries.unOp { + TODO() } - override fun StructureND.sum(): T = TODO("Not yet implemented") + override fun StructureND.sum(): T = unOp { + ops.sum(it, ops.constant(intArrayOf())) + }.value() - override fun StructureND.sum(dim: Int, keepDim: Boolean): Tensor { - TODO("Not yet implemented") + override fun StructureND.sum(dim: Int, keepDim: Boolean): TensorFlowOutput = unOp { + ops.sum(it, ops.constant(dim), Sum.keepDims(keepDim)) } - override fun StructureND.min(): T { - TODO("Not yet implemented") + override fun StructureND.min(): T = unOp { + ops.min(it, ops.constant(intArrayOf())) + }.value() + + override fun StructureND.min(dim: Int, keepDim: Boolean): Tensor = unOp { + ops.min(it, ops.constant(dim), Min.keepDims(keepDim)) } - override fun StructureND.min(dim: Int, keepDim: Boolean): Tensor { - TODO("Not yet implemented") + override fun StructureND.max(): T = unOp { + ops.max(it, ops.constant(intArrayOf())) + }.value() + + override fun StructureND.max(dim: Int, keepDim: Boolean): Tensor = unOp { + ops.max(it, ops.constant(dim), Max.keepDims(keepDim)) } - override fun StructureND.max(): T { - TODO("Not yet implemented") - } + override fun StructureND.argMax(dim: Int, keepDim: Boolean): Tensor = IntTensorFlowOutput( + graph, + ops.math.argMax(asTensorFlow().output, ops.constant(dim), TInt32::class.java).output() + ).actualTensor - override fun StructureND.max(dim: Int, keepDim: Boolean): Tensor { - TODO("Not yet implemented") - } - - override fun StructureND.argMax(dim: Int, keepDim: Boolean): Tensor { - TODO("Not yet implemented") - } + @OptIn(UnstableKMathAPI::class) + override fun export(arg: StructureND): StructureND = + if (arg is TensorFlowOutput) arg.actualTensor else arg } \ No newline at end of file diff --git a/kmath-tensorflow/src/test/kotlin/space/kscience/kmath/tensorflow/DoubleTensorFlowOps.kt b/kmath-tensorflow/src/test/kotlin/space/kscience/kmath/tensorflow/DoubleTensorFlowOps.kt new file mode 100644 index 000000000..b7a4b94b4 --- /dev/null +++ b/kmath-tensorflow/src/test/kotlin/space/kscience/kmath/tensorflow/DoubleTensorFlowOps.kt @@ -0,0 +1,19 @@ +package space.kscience.kmath.tensorflow + +import org.junit.jupiter.api.Test +import space.kscience.kmath.nd.StructureND +import space.kscience.kmath.nd.structureND +import space.kscience.kmath.operations.DoubleField + +class DoubleTensorFlowOps { + @Test + fun basicOps() { + val res = DoubleField.produceWithTF { + val initial = structureND(2, 2) { 1.0 } + + initial + (initial * 2.0) + } + println(StructureND.toString(res)) + } + +} \ No newline at end of file diff --git a/kmath-viktor/src/main/kotlin/space/kscience/kmath/viktor/ViktorFieldOpsND.kt b/kmath-viktor/src/main/kotlin/space/kscience/kmath/viktor/ViktorFieldOpsND.kt index e43bbbc6f..a6ab55624 100644 --- a/kmath-viktor/src/main/kotlin/space/kscience/kmath/viktor/ViktorFieldOpsND.kt +++ b/kmath-viktor/src/main/kotlin/space/kscience/kmath/viktor/ViktorFieldOpsND.kt @@ -3,6 +3,8 @@ * Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. */ +@file:OptIn(PerformancePitfall::class) + package space.kscience.kmath.viktor import org.jetbrains.bio.viktor.F64Array