<!--- @file:Suppress("ClassName") ---> <!--- @file:Suppress("PropertyName") --->

Kotlin∇: Type-safe Symbolic Differentiation for the JVM

Kotlin∇ is a type-safe automatic differentiation framework written in Kotlin. It allows users to express differentiable programs with higher-dimensional data structures and operators. We attempt to restrict syntactically valid constructions to those which are algebraically valid and can be checked at compile-time. By enforcing these constraints in the type system, it eliminates certain classes of runtime errors that may occur during the execution of a differentiable program. Due to type-inference, most type declarations may be safely omitted by the end-user. Kotlin∇ strives to be expressive, safe, and notationally similar to mathematics.

Table of contents

• Introduction
• Supported features
• Usage
  • Installation
  • Notation
  • Shape safety
  • Higher-rank
  • Higher-order
  • Example
  • Variable capture
• Visualization
  • Dataflow graphs
  • Plotting functions
  • Loss curves
• Testing and gradient checking
• How does it work?
  • Operator overloading
  • First-class functions
  • Multi-stage programming
  • Extension functions
  • Algebraic data types
  • Multiple dispatch
  • Shape-safe tensor operations
  • Intermediate representation
  • Property delegation
• Experimental ideas
  • Church encoding
  • Type classes
  • Type arithmetic
• Formal grammar
• UML diagram
• Comparison to other frameworks
• References
• Acknowledgements

Introduction

Inspired by Stalin∇, Autograd, DiffSharp, Myia, Nexus, Tangent, Lantern et al., Kotlin∇ attempts to port recent advancements in automatic differentiation (AD) to the Kotlin language. AD is useful for gradient descent and has a variety of applications in numerical optimization and machine learning. Our implementation adds a number of experimental ideas, including compile-time shape-safety, algebraic simplification and numerical stability checking with property-based testing. We aim to provide an algebraically-grounded implementation of AD for shape-safe tensor operations. Tensors in Kotlin∇ are represented as multidimensional arrays.

Features

Kotlin∇ currently supports the following features:

• Arithmetical operations on scalars, vectors and matrices
• Shape-safe vector and matrix algebra
• Partial and higher-order differentiation on scalars
• Property-based testing for numerical gradient checking
• Recovery of symbolic derivatives from AD

Additionally, it aims to support:

• PyTorch-style define-by-run semantics
• N-dimensional tensors and higher-order tensor operators
• Fully-general AD over control flow, variable reassignment (via delegation), and array programming, possibly using a typed IR such as Myia

All of these features are implemented without access to bytecode or special compiler tricks - just using higher-order functions and lambdas as shown in Lambda the Ultimate Backpropogator, embedded DSLs a la Lightweight Modular Staging, and ordinary generics. Please see below for a more detailed feature comparison.

Usage

Installation

Kotlin∇ is hosted on Maven Central. An example project is provided here.

Gradle

dependencies {
  implementation("ai.hypergraph:kotlingrad:0.4.7")
}

Maven

<dependency>
  <groupId>ai.hypergraph</groupId>
  <artifactId>kotlingrad</artifactId>
  <version>0.4.7</version>
</dependency>

Jupyter Notebook

To access Kotlin∇'s notebook support, use the following line magic:

@file:DependsOn("ai.hypergraph:kotlingrad:0.4.7")

For more information, explore the tutorial.

Notation

Kotlin∇ operators are higher-order functions, which take at most two inputs and return a single output, all of which are functions with the same numerical type, and whose shape is denoted using superscript in the rightmost column below.

| Math | Infix ^† | Prefix | Postfix^‡ | Operator Type Signature | |:------------------------------------------------------------------:|:-------------------------------:|:--------------------------------:|:-----------------------------------:|:-------------------------------------------------------------------------------:| | $$\mathbf{A}(\mathbf{B})$$<br>$$\mathbf{A}\circ\mathbf{B}$$ | a(b)<br>a of b | | | $$(\texttt{a}: ℝ^{τ}→ℝ^{π}, \texttt{b}: ℝ^{λ} → ℝ^{τ}) → (ℝ^{λ}→ℝ^{π})$$ | | $$\mathbf{A}\pm\mathbf{B}$$ | a + b<br>a - b | plus(a, b)<br>minus(a, b) | | $$(\texttt{a}: ℝ^{τ}→ℝ^{π}, \texttt{b}: ℝ^{λ} → ℝ^{π}) → (ℝ^{?}→ℝ^{π})$$ | | $$\mathbf{A}\mathbf{B}$$ | a * b<br>a.times(b) | times(a, b) | | $$(\texttt{a}: ℝ^{τ}→ℝ^{m×n}, \texttt{b}: ℝ^{λ}→ℝ^{n×p}) → (ℝ^{?}→ℝ^{m×p})$$ | | $$\frac{\mathbf{A}}{\mathbf{B}}$$<br>$$\mathbf{A}\mathbf{B}^{-1}$$ | a / b<br>a.div(b) | div(a, b) | | $$(\texttt{a}: ℝ^{τ}→ℝ^{m×n}, \texttt{b}: ℝ^{λ}→ℝ^{p×n}) → (ℝ^{?}→ℝ^{m×p})$$ | | $$\pm\mathbf{A}$$ | | -a<br>+a | a.neg()<br>a.pos() | $$(\texttt{a}: ℝ^{τ}→ℝ^{π}) → (ℝ^{τ}→ℝ^{π})$$ | | $$\sin{a}$$<br>$$\cos{a}$$<br>$$\tan{a}$$ | | sin(a)<br>cos(a)<br>tan(a) | a.sin()<br>a.cos()<br>a.tan() | $$(\texttt{a}: ℝ→ℝ) → (ℝ→ℝ)$$ | | $$\ln{a}$$ | | ln(a)<br>log(a) | a.ln()<br>a.log() | $$(\texttt{a}: ℝ^{τ}→ℝ^{m×m}) → (ℝ^{τ}→ℝ^{m×m})$$ | | $$\log_{b}a$$ | a.log(b) | log(a, b) | | $$(\texttt{a}: ℝ^{τ}→ℝ^{m×m}, \texttt{b}: ℝ^{λ}→ℝ^{m×m}) → (ℝ^{?}→ℝ)$$ | | $$\mathbf{A}^b$$ | a.pow(b) | pow(a, b) | | $$(\texttt{a}: ℝ^{τ}→ℝ^{m×m}, \texttt{b}: ℝ^{λ}→ℝ) → (ℝ^{?}→ℝ^{m×m})$$ | | $$\sqrt{A}$$<br>$$\sqrt[3]{A}$$ | a.pow(1.0/2)<br>a.root(3) | sqrt(a)<br>cbrt(a) | a.sqrt()<br>a.cbrt() | $$(\texttt{a}: ℝ^{τ}→ℝ^{m×m}) → (ℝ^{τ}→ℝ^{m×m})$$ | | $$\frac{da}{db},\frac{\partial{a}}{\partial{b}}$$ <br> $$D_b{a}$$ | a.d(b)<br>d(a) / d(b) | grad(a)[b] | | $$(\texttt{a}: C(ℝ^{τ}→ℝ)^{*}, \texttt{b}: C(ℝ^{λ}→ℝ)) → (ℝ^{?}→ℝ)$$ | | $$\nabla{a}$$ | | grad(a) | `a.