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.

Introduction
Supported features
Usage
Visualization
Testing and gradient checking
How does it work?
Experimental ideas
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})$$ $$\mathbf{A}\circ\mathbf{B}$$	`a(b)` `a of b`			$$(\texttt{a}: ℝ^{τ}→ℝ^{π}, \texttt{b}: ℝ^{λ} → ℝ^{τ}) → (ℝ^{λ}→ℝ^{π})$$
$$\mathbf{A}\pm\mathbf{B}$$	`a + b` `a - b`	`plus(a, b)` `minus(a, b)`		$$(\texttt{a}: ℝ^{τ}→ℝ^{π}, \texttt{b}: ℝ^{λ} → ℝ^{π}) → (ℝ^{?}→ℝ^{π})$$
$$\mathbf{A}\mathbf{B}$$	`a * b` `a.times(b)`	`times(a, b)`		$$(\texttt{a}: ℝ^{τ}→ℝ^{m×n}, \texttt{b}: ℝ^{λ}→ℝ^{n×p}) → (ℝ^{?}→ℝ^{m×p})$$
$$\frac{\mathbf{A}}{\mathbf{B}}$$ $$\mathbf{A}\mathbf{B}^{-1}$$	`a / b` `a.div(b)`	`div(a, b)`		$$(\texttt{a}: ℝ^{τ}→ℝ^{m×n}, \texttt{b}: ℝ^{λ}→ℝ^{p×n}) → (ℝ^{?}→ℝ^{m×p})$$
$$\pm\mathbf{A}$$		`-a` `+a`	`a.neg()` `a.pos()`	$$(\texttt{a}: ℝ^{τ}→ℝ^{π}) → (ℝ^{τ}→ℝ^{π})$$
$$\sin{a}$$ $$\cos{a}$$ $$\tan{a}$$		`sin(a)` `cos(a)` `tan(a)`	`a.sin()` `a.cos()` `a.tan()`	$$(\texttt{a}: ℝ→ℝ) → (ℝ→ℝ)$$
$$\ln{a}$$		`ln(a)` `log(a)`	`a.ln()` `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}$$ $$\sqrt[3]{A}$$	`a.pow(1.0/2)` `a.root(3)`	`sqrt(a)` `cbrt(a)`	`a.sqrt()` `a.cbrt()`	$$(\texttt{a}: ℝ^{τ}→ℝ^{m×m}) → (ℝ^{τ}→ℝ^{m×m})$$
$$\frac{da}{db},\frac{\partial{a}}{\partial{b}}$$ $$D_b{a}$$	`a.d(b)` `d(a) / d(b)`	`grad(a)[b]`		$$(\texttt{a}: C(ℝ^{τ}→ℝ)^{*}, \texttt{b}: C(ℝ^{λ}→ℝ)) → (ℝ^{?}→ℝ)$$
$$\nabla{a}$$		`grad(a)`	`a.grad()`	$$(\texttt{a}: C(ℝ^{τ}→ℝ)) → (ℝ^{τ}→ℝ^{τ})$$
$$\nabla_{\mathbf{B}}a$$	`a.d(b)` `a.grad(b)`	`grad(a, b)` `grad(a)[b]`		$$(\texttt{a}: C(ℝ^{τ}→ℝ^{π}), \texttt{b}: C(ℝ^{λ}→ℝ^{ω})) → (ℝ^{?}→ℝ^{π×ω})$$
$$\nabla\cdot{\mathbf{A}}$$		`divg(a)`	`a.divg()`	$$(\texttt{a}: C(ℝ^{τ}→ℝ^{m})) → (ℝ^{τ}→ℝ)$$
$$\nabla\times{\mathbf{A}}$$		`curl(a)`	`a.curl()`	$$(\texttt{a}: C(ℝ^{3}→ℝ^{3})) → (ℝ^{3}→ℝ^{3})$$
$$\mathcal{J}(\mathbf{A})$$		`grad(a)`	`a.grad()`	$$(\texttt{a}: C(ℝ^{τ}→ℝ^{m})) → (ℝ^{τ}→ℝ^{m×τ})$$
$$\mathbf{H}(a)$$		`hess(a)`	`a.hess()`	$$(\texttt{a}: C(ℝ^{τ}→ℝ)) → (ℝ^{τ}→ℝ^{τ×τ})$$
$$\Delta{a},\nabla^{2}a$$		`lapl(a)`	`a.lapl()`	$$(\texttt{a}: C(ℝ^{τ}→ℝ)) → (ℝ^{τ}→ℝ^{τ})$$

ℝ can be a Double, Float or BigDecimal. Specialized operators are defined for subsets of ℝ, e.g., Int, Short or BigInteger for subsets of ℤ, however differentiation is only defined for continuously differentiable functions on ℝ.

^† a and b are higher-order functions. These may be constants (e.g., 0, 1.0), variables (e.g., Var()) or expressions (e.g., x + 1, 2 * x + y).

^‡ For infix notation, . is optional. Parentheses are also optional depending on precedence.

^§ Matrix division is defined iff B is invertible, although it could be possible to redefine this operator using the Moore-Penrose inverse.

^∗ Where C(ℝ^m) is the space of all continuous functions over ℝ. If the function is not over ℝ, it will fail at compile-time. If the function is over ℝ but not continuous differentiable at the point under consideration, it will fail at runtime.

^? The input shape is tracked at runtime, but not at the type level. While it would be nice to infer a union type bound over the inputs of binary functions, it is likely impossible using the Kotlin type system without great effort. If the user desires type checking when invoking higher order functions with literal values, they will need to specify the combined input type explicitly or do so at runtime.

^{τ, λ, π, ω} Arbitrary products.

Higher-Rank Derivatives

Kotlin∇ supports derivatives between tensors of up to rank 2. The shape of a tensor derivative depends on (1) the shape of the function under differentiation and (2) the shape of the variable with respect to which we are differentiating.