AdOptimize

.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: |

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<org.jetbrains.kotlin.gradle.tasks.KotlinCompile>() {
 publishing {
     publications {
         create<MavenPublication>("maven") {
-            groupId = "org.diffkt"
+            groupId = "org.diffkt.adopt"
             artifactId = "api"
             version = project.version.toString()
             from(components["java"])

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) -> Pair<DTensor, (DTensor)->DTensor>
+ */
+annotation class PrimalAndPullback
+
+@StackImpl
+class CodeGenStack<T> {
+    val data = arrayListOf<T>()
+    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()
+}

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 {
 
     /**

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<DTensor> {
 

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

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")

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

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
 }

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)
 }

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 })

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

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

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"

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<KotlinCompile>() {
-        kotlinOptions.jvmTarget = "11"
+        kotlinOptions.jvmTarget = "1.8"
         kotlinOptions.freeCompilerArgs += "-XXLanguage:+ProperCheckAnnotationsTargetInTypeUsePositions"
     }