diff --git a/.github/workflows/config.yml b/.github/workflows/config.yml index 345277ac..285894e4 100644 --- a/.github/workflows/config.yml +++ b/.github/workflows/config.yml @@ -13,10 +13,10 @@ jobs: contents: read packages: write steps: - - name: set up jdk 11 + - name: set up jdk 8 uses: actions/setup-java@v2 with: - java-version: '11' + java-version: '8' distribution: 'adopt' - name: install diffkt dependencies run: | diff --git a/kotlin/api/build.gradle.kts b/kotlin/api/build.gradle.kts index e079abf7..7d03dd13 100644 --- a/kotlin/api/build.gradle.kts +++ b/kotlin/api/build.gradle.kts @@ -9,6 +9,7 @@ import org.jetbrains.dokka.gradle.DokkaTask plugins { `maven-publish` + id("meta-diffkt-differentiable-api-preprocessor") version "0.0.1.3" id("shapeKt") version "1.0" id("org.jetbrains.dokka") version "1.6.0" } @@ -28,6 +29,23 @@ repositories { mavenCentral() } +differentiableApiPreprocessor { + this.stackImplAnnotation("org.diffkt.adOptimize.StackImpl") + this.boxedPrimitive("org.diffkt.adOptimize.BoxedPrimitive") + this.scalarRoot("org.diffkt.adOptimize.ScalarRoot") + this.primalAndPullbackAnnotation("org.diffkt.adOptimize.PrimalAndPullback") + this.reverseAnnotation("org.diffkt.adOptimize.ReverseDifferentiable") + this.unboxedFunction("org.diffkt.adOptimize.ToUnboxedFunction") + val userDir = System.getProperty("user.dir") + val pathToResources = "$userDir/api/src/main/resources" + this.resourcesPath(pathToResources) + this.toReverseAnnotation("org.diffkt.adOptimize.ToReverse") + this.dTensorAnnotation("org.diffkt.adOptimize.DTensorRoot") + this.reverseScalarOperationsAnnotation("org.diffkt.adOptimize.ReverseScalarOperations") + this.scalarNoop("org.diffkt.adOptimize.ScalarNoop") + this.forwardDifferentiable("org.diffkt.adOptimize.ForwardDifferentiable") +} + dependencies { implementation(group = "net.bytebuddy", name = "byte-buddy", version="1.12.7") compileOnly("shapeKt:annotations:1.0") @@ -63,7 +81,7 @@ tasks.withType() { publishing { publications { create("maven") { - groupId = "org.diffkt" + groupId = "org.diffkt.adopt" artifactId = "api" version = project.version.toString() from(components["java"]) diff --git a/kotlin/api/src/main/kotlin/org/diffkt/ADConfig.kt b/kotlin/api/src/main/kotlin/org/diffkt/ADConfig.kt new file mode 100644 index 00000000..02c0d96e --- /dev/null +++ b/kotlin/api/src/main/kotlin/org/diffkt/ADConfig.kt @@ -0,0 +1,51 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +package org.diffkt.adOptimize + +annotation class StackImpl +annotation class BoxedPrimitive(val value:String) +annotation class ScalarRoot +annotation class ToUnboxedFunction(val functionName:String) +annotation class ToReverse(val fqClass:String) +annotation class DTensorRoot +annotation class ReverseScalarOperations +annotation class ForwardDifferentiable(val tangentProperty: String) + +/** + * a function is marked as a scalar noop if it accepts exactly one potentially + * active operand and returns the same type. The return value MUST be the implicit receiver + */ +annotation class ScalarNoop + +/** + * The reverse differentiable scalar type should be annotated with this annotation + * so the compiler plugin can build implementations of reverse nodes + */ +annotation class ReverseDifferentiable(val primalField:String, val upstreamField:String, val backpropogateMethod:String, val pushbackMethod:String, val derivativeID:String) + +/** + * This is used by the compiler when it cannot inline a function call. + * The compiler expects the signature of the function annotated with this annotation + * to be (DTensor, (DTensor) -> DTensor) -> PairDTensor> + */ +annotation class PrimalAndPullback + +@StackImpl +class CodeGenStack { + val data = arrayListOf() + fun push(d:T){ + data.add(d) + } + fun pop():T { + val x = data.last() + data.removeLast() + return x + } + fun top():T = data.last() + fun notEmpty() = data.isNotEmpty() +} \ No newline at end of file diff --git a/kotlin/api/src/main/kotlin/org/diffkt/DScalar.kt b/kotlin/api/src/main/kotlin/org/diffkt/DScalar.kt index 910335d6..63bda363 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/DScalar.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/DScalar.kt @@ -8,6 +8,7 @@ package org.diffkt import shapeTyping.annotations.SType +import org.diffkt.adOptimize.ScalarRoot /** * A differentiable scalar (float). @@ -16,6 +17,7 @@ import shapeTyping.annotations.SType * - a [ForwardScalar] for forward differentiation with a [DScalar] primal value, and a [DTensor] tangent, or * - a [ReverseScalar] for reverse mode differentiation. */ +@ScalarRoot interface DScalar : @SType("[]") DTensor { /** diff --git a/kotlin/api/src/main/kotlin/org/diffkt/DTensor.kt b/kotlin/api/src/main/kotlin/org/diffkt/DTensor.kt index f7f468ba..32fc4123 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/DTensor.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/DTensor.kt @@ -8,6 +8,7 @@ package org.diffkt import shapeTyping.annotations.SType +import org.diffkt.adOptimize.DTensorRoot /** * Interface for a differentiable tensor. @@ -17,6 +18,7 @@ import shapeTyping.annotations.SType * - a [ForwardTensor] for forward differentiation, or * - a [ReverseTensor] for reverse mode differentiation. */ +@DTensorRoot @SType("S: Shape") interface DTensor: Differentiable { diff --git a/kotlin/api/src/main/kotlin/org/diffkt/FloatScalar.kt b/kotlin/api/src/main/kotlin/org/diffkt/FloatScalar.kt index b4c1b5be..c5a72587 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/FloatScalar.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/FloatScalar.kt @@ -7,12 +7,15 @@ package org.diffkt +import org.diffkt.adOptimize.BoxedPrimitive + /** * A differentiable tensor of rank 0 containing a single float (a FloatTensor wrapper around a float). * * @property value The floating point number to wrap with a FloatTensor. * @constructor Creates a FloatScalar initialized to value. */ +@BoxedPrimitive("value") class FloatScalar(val value: Float) : FloatTensor(), DScalar { override val derivativeID: DerivativeID get() = NoDerivativeID override val operations: Operations diff --git a/kotlin/api/src/main/kotlin/org/diffkt/Power.kt b/kotlin/api/src/main/kotlin/org/diffkt/Power.kt index b9e5e511..7e5b412c 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/Power.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/Power.kt @@ -7,6 +7,7 @@ package org.diffkt +import org.diffkt.adOptimize.ToUnboxedFunction import shapeTyping.annotations.SType // Tensor powers @@ -17,10 +18,12 @@ fun DTensor.pow(x: DScalar): DTensor = exp(x * ln(this)) // a^b == e^(b ln a) fun DScalar.pow(x: DScalar): DScalar = exp(x * ln(this)) +@ToUnboxedFunction("kotlin.math.pow") fun DScalar.pow(x: Float): DScalar = (this as DTensor).pow(x) as DScalar fun DScalar.pow(x: Int): DScalar = (this as DTensor).pow(x.toFloat()) as DScalar +@ToUnboxedFunction("kotlin.math.pow") fun DTensor.pow(x: Int): DTensor = this.pow(x.toFloat()) @SType("S: Shape") diff --git a/kotlin/api/src/main/kotlin/org/diffkt/ReverseDerivative.kt b/kotlin/api/src/main/kotlin/org/diffkt/ReverseDerivative.kt index 0ef6d943..54e50c5c 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/ReverseDerivative.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/ReverseDerivative.kt @@ -10,6 +10,7 @@ package org.diffkt import org.diffkt.reverse.ReverseDerivativeID import org.diffkt.reverse.ReverseScalar import org.diffkt.reverse.ReverseTensor +import org.diffkt.adOptimize.PrimalAndPullback // ********** General Reverse Derivative of Univariate Functions ********** @@ -54,6 +55,7 @@ fun vjp( ) = primalAndVjp(x, vf, f).second // Also known as VJP (vector-Jacobian product) +@PrimalAndPullback internal fun primalAndPullback( x: DTensor, f: (x: DTensor) -> DTensor diff --git a/kotlin/api/src/main/kotlin/org/diffkt/Sigmoid.kt b/kotlin/api/src/main/kotlin/org/diffkt/Sigmoid.kt index 48fae59f..55afda8b 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/Sigmoid.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/Sigmoid.kt @@ -9,6 +9,7 @@ package org.diffkt import kotlin.math.exp import shapeTyping.annotations.SType +import org.diffkt.adOptimize.ToUnboxedFunction /** * Compute the sigmoid for a single floating-point value. @@ -23,6 +24,7 @@ internal fun sigmoidElem(x: Float): Float { } } +@ToUnboxedFunction("org.diffkt.sigmoidElem") fun sigmoid(x: DScalar): DScalar { return x.operations.sigmoid(x) as DScalar } diff --git a/kotlin/api/src/main/kotlin/org/diffkt/Transcendental.kt b/kotlin/api/src/main/kotlin/org/diffkt/Transcendental.kt index f468437a..9f477ad8 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/Transcendental.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/Transcendental.kt @@ -7,23 +7,29 @@ package org.diffkt +import org.diffkt.adOptimize.ToUnboxedFunction + // ========== // Transcendentals // ========== // Trig +@ToUnboxedFunction("kotlin.math.sin") fun sin(x: DScalar): DScalar { return x.operations.sin(x) as DScalar } +@ToUnboxedFunction("kotlin.math.sin") fun sin(x: DTensor): DTensor { return x.operations.sin(x) } +@ToUnboxedFunction("kotlin.math.cos") fun cos(x: DScalar): DScalar { return x.operations.cos(x) as DScalar } +@ToUnboxedFunction("kotlin.math.cos") fun cos(x: DTensor): DTensor { return x.operations.cos(x) } @@ -45,6 +51,7 @@ fun atan(x: DTensor): DTensor { return x.operations.atan(x) } +@ToUnboxedFunction("kotlin.math.exp") fun exp(x: DScalar): DScalar { return x.operations.exp(x) as DScalar } @@ -53,10 +60,12 @@ fun exp(x: DTensor): DTensor { return x.operations.exp(x) } +@ToUnboxedFunction("kotlin.math.ln") fun ln(x: DScalar): DScalar { return x.operations.ln(x) as DScalar } +@ToUnboxedFunction("kotlin.math.ln") fun ln(x: DTensor): DTensor { return x.operations.ln(x) } diff --git a/kotlin/api/src/main/kotlin/org/diffkt/Utils.kt b/kotlin/api/src/main/kotlin/org/diffkt/Utils.kt index d2bd9406..7acc9fb9 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/Utils.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/Utils.kt @@ -9,6 +9,7 @@ package org.diffkt import org.diffkt.external.ExternalLib import org.diffkt.reverse.ReverseTensor +import org.diffkt.adOptimize.ScalarNoop internal fun identityGradientofSameKind(x: DTensor, halfShape: Shape = x.shape): DTensor { return x.operations.identityGradientOfSameKind(x, halfShape) @@ -82,6 +83,7 @@ val IntArray.product get() = run { product } +@ScalarNoop fun DTensor.expandToTangent(tangent: DTensor): DTensor { if (this.shape == tangent.shape) return this val ones = Shape(IntArray(tangent.shape.rank - this.shape.rank) { 1 }) diff --git a/kotlin/api/src/main/kotlin/org/diffkt/forward/ForwardScalar.kt b/kotlin/api/src/main/kotlin/org/diffkt/forward/ForwardScalar.kt index e9c4e38e..4e4c4e6c 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/forward/ForwardScalar.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/forward/ForwardScalar.kt @@ -10,10 +10,12 @@ package org.diffkt.forward import org.diffkt.DScalar import org.diffkt.DTensor import org.diffkt.Operations +import org.diffkt.adOptimize.ForwardDifferentiable /** * A differentiable dual scalar (for forward derivatives) */ +@ForwardDifferentiable("tangent") open class ForwardScalar protected constructor( primal: DScalar, derivativeID: ForwardDerivativeID diff --git a/kotlin/api/src/main/kotlin/org/diffkt/reverse/ReverseScalar.kt b/kotlin/api/src/main/kotlin/org/diffkt/reverse/ReverseScalar.kt index deabb098..b5daaa33 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/reverse/ReverseScalar.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/reverse/ReverseScalar.kt @@ -10,10 +10,12 @@ package org.diffkt.reverse import org.diffkt.DScalar import org.diffkt.gpu.GpuFloatScalar import org.diffkt.Operations +import org.diffkt.adOptimize.ReverseDifferentiable /** * A scalar for reverse mode differentiation. */ +@ReverseDifferentiable("primal", "upstream", "backpropagate", "pushback", "derivativeID") abstract class ReverseScalar(override val primal: DScalar, derivativeID: ReverseDerivativeID) : ReverseTensor(primal, derivativeID), DScalar { override val operations: Operations diff --git a/kotlin/api/src/main/kotlin/org/diffkt/reverse/ReverseScalarOperationsImpl.kt b/kotlin/api/src/main/kotlin/org/diffkt/reverse/ReverseScalarOperationsImpl.kt index 41379a7c..7f747df4 100644 --- a/kotlin/api/src/main/kotlin/org/diffkt/reverse/ReverseScalarOperationsImpl.kt +++ b/kotlin/api/src/main/kotlin/org/diffkt/reverse/ReverseScalarOperationsImpl.kt @@ -11,7 +11,9 @@ import org.diffkt.* import org.diffkt.model.BatchNormResult import org.diffkt.random.RandomKey import shapeTyping.annotations.AllowUnreduced +import org.diffkt.adOptimize.ReverseScalarOperations +@ReverseScalarOperations @AllowUnreduced internal open class ReverseScalarOperationsImpl: Operations { override val name get() = "ReverseScalar" diff --git a/kotlin/build.gradle.kts b/kotlin/build.gradle.kts index 3bf54f75..0f8c3272 100644 --- a/kotlin/build.gradle.kts +++ b/kotlin/build.gradle.kts @@ -26,12 +26,12 @@ allprojects { apply(plugin = "java") java { - sourceCompatibility = JavaVersion.VERSION_11 - targetCompatibility = JavaVersion.VERSION_11 + sourceCompatibility = JavaVersion.VERSION_1_8 + targetCompatibility = JavaVersion.VERSION_1_8 } tasks.withType() { - kotlinOptions.jvmTarget = "11" + kotlinOptions.jvmTarget = "1.8" kotlinOptions.freeCompilerArgs += "-XXLanguage:+ProperCheckAnnotationsTargetInTypeUsePositions" }