Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Upstreaming Swift AutoDiff

Upstreaming Swift AutoDiff

Author: Richard Wei (rxwei@google.com) on behalf of the Swift for TensorFlow team

Last updated: October 2, 2019

Overview

The differentiable programming feature (AutoDiff) has been incubated in the 'tensorflow' branch of apple/swift since December 2017 and released as part of the Swift for TensorFlow toolchains. The Differentiable Programming Mega-Proposal, which serves as a manifesto, received general positive feedback from the community, but there is a long way between receiving conceptual approval and obtaining Swift Evolution approval of such a large feature. We would like to merge the pieces into the 'master' branch under a gate to further development and bake the feature on master, just like Apple develops its major features.

Motivation

The maturity of the differentiable programming feature is improving and we would like to start a series of formal reviews on Swift Evolution. We realize that these will go smoother if the code is formally reviewed, and the best way to do that is to merge the code into the 'master' branch while disabling user-visible changes under a feature gate. As portions of the feature are approved by the Swift Evolution process, they can be moved outside of the protection of the feature gate.

Feature gating

The user-visible part of the differentiable programming feature, including compiler changes and standard library changes, will be protected under an -enable-experimental-differentiable-programming frontend flag. Differentiation APIs, including top-level functions that allow users to take the derivative of a function, will be located in a _Differentiation module.

Upstreaming the differentiable programming feature will not change the compiler’s behavior when -enable-experimental-differentiable-programming is not provided to the frontend, except for some syntax unavailability diagnostics when a differentiation-specific attribute or protocol is used without the -enable-experimental-differentiable-programming flag.

Components

The differentiable programming feature requires changes to many parts of the core language. The following graph demonstrates the dependencies and roughly how things should be upstreamed piece-by-piece.

Components graph

Testing

Each upstreaming patch will follow existing testing requirements in apple/swift. Before differentiable programming becomes accepted via Swift Evolution, all differentiable programming tests will be located under directory 'test/AutoDiff/' in apple/swift. A DifferentiationUnittest module will be created in 'stdlib/private/' to contain lifetime-tracked differentiable data structures that differentiable programming runtime unit tests use. A fuzzing script for combining language features and different math that generates source code for testing will be upstreamed.

Responsibility

Code owners of the differentiable programming feature include Richard Wei (@rxwei), Dan Zheng (@dan-zheng), and Marc Rasi (@marcrasi). Upon any breakages of differentiable programming tests caused by changes in the 'master' branch, the code owners will aim to respond within a business day after any code owner is tagged on GitHub. Actual engineering fixes will be assessed collaboratively on a case-by-case basis.

Evolution reviews

The Differentiable Programming Mega-Proposal will be repositioned as a manifesto and be upstreamed to the 'docs/' directory in apple/swift. The authors will draft smaller proposals that contain references to the manifesto for detailed design and math definitions. The smaller proposals will contain significantly more code examples describing the additions and how to use the feature. The first such proposal will include minimal changes that make the basic differentiable programming work, including the following user-visible additions:

  • The Differentiable protocol
  • @differentiable, @differentiating and @transposing declaration attributes
  • The @differentiable function type attribute
  • Basic differentiation APIs: valueWithDerivative(at:in:) and transpose(of:)

Later pitches and proposals may include components like derived conformances for AdditiveArithmetic and Differentiable, more differentiation APIs like gradient(of:), differentiable higher-order functions, and higher-order differentiation. These topics will be out of scope for the initial pitch.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment