[AutoDiff upstream] Upstream _Differentiation APIs.

stdlib/public/Differentiation/AnyDerivative.swift

@@ -0,0 +1,292 @@
+//===--- AnyDerivative.swift ----------------------------------*- swift -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2019 - 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines `AnyDerivative`, a type-erased wrapper for
+// `Differentiable.TangentVector` associated type implementations.
+//
+//===----------------------------------------------------------------------===//
+
+import Swift
+
+@usableFromInline
+internal protocol _AnyDerivativeBox {
+  // `Equatable` requirements (implied by `AdditiveArithmetic`).
+  func _isEqual(to other: _AnyDerivativeBox) -> Bool
+  func _isNotEqual(to other: _AnyDerivativeBox) -> Bool
+
+  // `AdditiveArithmetic` requirements.
+  static var _zero: _AnyDerivativeBox { get }
+  func _adding(_ x: _AnyDerivativeBox) -> _AnyDerivativeBox
+  func _subtracting(_ x: _AnyDerivativeBox) -> _AnyDerivativeBox
+
+  // `Differentiable` requirements.
+  mutating func _move(along direction: _AnyDerivativeBox)
+
+  /// The underlying base value, type-erased to `Any`.
+  var _typeErasedBase: Any { get }
+
+  /// Returns the underlying value unboxed to the given type, if possible.
+  func _unboxed<U>(to type: U.Type) -> U?
+  where U: Differentiable, U.TangentVector == U
+}
+
+extension _AnyDerivativeBox {
+  /// Returns true if the underlying value has type `AnyDerivative.OpaqueZero`.
+  @inlinable
+  func _isOpaqueZero() -> Bool {
+    return _unboxed(to: AnyDerivative.OpaqueZero.self) != nil
+  }
+}
+
+@inline(never)
+@usableFromInline
+internal func _derivativeTypeMismatch(
+  _ x: Any.Type, _ y: Any.Type, file: StaticString = #file, line: UInt = #line
+) -> Never {
+  preconditionFailure(
+    """
+    Derivative type mismatch: \
+    \(String(reflecting: x)) and \(String(reflecting: y))
+    """, file: file, line: line)
+}
+
+@frozen
+@usableFromInline
+internal struct _ConcreteDerivativeBox<T>: _AnyDerivativeBox
+where T: Differentiable, T.TangentVector == T {
+  /// The underlying base value.
+  @usableFromInline
+  var _base: T
+
+  @inlinable
+  internal init(_ base: T) {
+    self._base = base
+  }
+
+  /// The underlying base value, type-erased to `Any`.
+  @inlinable
+  var _typeErasedBase: Any {
+    return _base
+  }
+
+  @inlinable
+  func _unboxed<U>(to type: U.Type) -> U?
+  where U: Differentiable, U.TangentVector == U {
+    return (self as? _ConcreteDerivativeBox<U>)?._base
+  }
+
+  // `Equatable` requirements (implied by `AdditiveArithmetic`).
+  @inlinable
+  func _isEqual(to other: _AnyDerivativeBox) -> Bool {
+    return _base == other._unboxed(to: T.self)
+  }
+  @inlinable
+  func _isNotEqual(to other: _AnyDerivativeBox) -> Bool {
+    return _base != other._unboxed(to: T.self)
+  }
+
+  // `AdditiveArithmetic` requirements.
+
+  @inlinable
+  static var _zero: _AnyDerivativeBox {
+    return _ConcreteDerivativeBox(T.zero)
+  }
+
+  @inlinable
+  func _adding(_ x: _AnyDerivativeBox) -> _AnyDerivativeBox {
+    // 0 + x = x
+    if _isOpaqueZero() {
+      return x
+    }
+    // y + 0 = y
+    if x._isOpaqueZero() {
+      return self
+    }
+    guard let xBase = x._unboxed(to: T.self) else {
+      _derivativeTypeMismatch(T.self, type(of: x._typeErasedBase))
+    }
+    return _ConcreteDerivativeBox(_base + xBase)
+  }
+
+  @inlinable
+  func _subtracting(_ x: _AnyDerivativeBox) -> _AnyDerivativeBox {
+    // y - 0 = y
+    if x._isOpaqueZero() {
+      return self
+    }
+    // 0 - x = -x
+    if _isOpaqueZero() {
+      return type(of: x)._zero._subtracting(x)
+    }
+    guard let xBase = x._unboxed(to: T.self) else {
+      _derivativeTypeMismatch(T.self, type(of: x._typeErasedBase))
+    }
+    return _ConcreteDerivativeBox(_base - xBase)
+  }
+
+  // `Differentiable` requirements.
+  @inlinable
+  mutating func _move(along direction: _AnyDerivativeBox) {
+    if direction._isOpaqueZero() {
+      return
+    }
+    // The case where `self._isOpaqueZero()` returns true is handled in
+    // `AnyDerivative.move(along:)`.
+    guard
+      let directionBase =
+        direction._unboxed(to: T.TangentVector.self)
+    else {
+      _derivativeTypeMismatch(T.self, type(of: direction._typeErasedBase))
+    }
+    _base.move(along: directionBase)
+  }
+}
+
+/// A type-erased derivative value.
+///
+/// The `AnyDerivative` type forwards its operations to an arbitrary underlying
+/// base derivative value conforming to `Differentiable` and
+/// `AdditiveArithmetic`, hiding the specifics of the underlying value.
+@frozen
+public struct AnyDerivative: Differentiable & AdditiveArithmetic {
+  @usableFromInline
+  internal var _box: _AnyDerivativeBox
+
+  @inlinable
+  internal init(_box: _AnyDerivativeBox) {
+    self._box = _box
+  }
+
+  /// The underlying base value.
+  @inlinable
+  public var base: Any {
+    return _box._typeErasedBase
+  }
+
+  /// Creates a type-erased derivative from the given derivative.
+  @inlinable
+  @differentiable
+  public init<T>(_ base: T) where T: Differentiable, T.TangentVector == T {
+    self._box = _ConcreteDerivativeBox<T>(base)
+  }
+
+  @inlinable
+  @derivative(of: init)
+  internal static func _vjpInit<T>(
+    _ base: T
+  ) -> (value: AnyDerivative, pullback: (AnyDerivative) -> T.TangentVector)
+  where T: Differentiable, T.TangentVector == T {
+    return (AnyDerivative(base), { v in v.base as! T.TangentVector })
+  }
+
+  @inlinable
+  @derivative(of: init)
+  internal static func _jvpInit<T>(
+    _ base: T
+  ) -> (value: AnyDerivative, differential: (T.TangentVector) -> AnyDerivative)
+  where T: Differentiable, T.TangentVector == T {
+    return (AnyDerivative(base), { dbase in AnyDerivative(dbase) })
+  }
+
+  public typealias TangentVector = AnyDerivative
+
+  // `Equatable` requirements (implied by `AdditiveArithmetic`).
+  @inlinable
+  public static func == (lhs: AnyDerivative, rhs: AnyDerivative) -> Bool {
+    return lhs._box._isEqual(to: rhs._box)
+  }
+  @inlinable
+  public static func != (lhs: AnyDerivative, rhs: AnyDerivative) -> Bool {
+    return lhs._box._isNotEqual(to: rhs._box)
+  }
+
+  // `AdditiveArithmetic` requirements.
+
+  /// Internal struct representing an opaque zero value.
+  @frozen
+  @usableFromInline
+  internal struct OpaqueZero: Differentiable & AdditiveArithmetic {}
+
+  @inlinable
+  public static var zero: AnyDerivative {
+    return AnyDerivative(
+      _box: _ConcreteDerivativeBox<OpaqueZero>(OpaqueZero.zero))
+  }
+
+  @inlinable
+  public static func + (
+    lhs: AnyDerivative, rhs: AnyDerivative
+  ) -> AnyDerivative {
+    return AnyDerivative(_box: lhs._box._adding(rhs._box))
+  }
+
+  @derivative(of: +)
+  @inlinable
+  internal static func _vjpAdd(
+    lhs: AnyDerivative, rhs: AnyDerivative
+  ) -> (
+    value: AnyDerivative,
+    pullback: (AnyDerivative) -> (AnyDerivative, AnyDerivative)
+  ) {
+    return (lhs + rhs, { v in (v, v) })
+  }
+
+  @derivative(of: +)
+  @inlinable
+  internal static func _jvpAdd(
+    lhs: AnyDerivative, rhs: AnyDerivative
+  ) -> (
+    value: AnyDerivative,
+    differential: (AnyDerivative, AnyDerivative) -> (AnyDerivative)
+  ) {
+    return (lhs + rhs, { (dlhs, drhs) in dlhs + drhs })
+  }
+
+  @inlinable
+  public static func - (
+    lhs: AnyDerivative, rhs: AnyDerivative
+  ) -> AnyDerivative {
+    return AnyDerivative(_box: lhs._box._subtracting(rhs._box))
+  }
+
+  @derivative(of: -)
+  @inlinable
+  internal static func _vjpSubtract(
+    lhs: AnyDerivative, rhs: AnyDerivative
+  ) -> (
+    value: AnyDerivative,
+    pullback: (AnyDerivative) -> (AnyDerivative, AnyDerivative)
+  ) {
+    return (lhs - rhs, { v in (v, .zero - v) })
+  }
+
+  @derivative(of: -)
+  @inlinable
+  internal static func _jvpSubtract(
+    lhs: AnyDerivative, rhs: AnyDerivative
+  ) -> (
+    value: AnyDerivative,
+    differential: (AnyDerivative, AnyDerivative) -> AnyDerivative
+  ) {
+    return (lhs - rhs, { (dlhs, drhs) in dlhs - drhs })
+  }
+
+  // `Differentiable` requirements.
+  @inlinable
+  public mutating func move(along direction: TangentVector) {
+    if _box._isOpaqueZero() {
+      _box = direction._box
+      return
+    }
+    _box._move(along: direction._box)
+  }
+}

stdlib/public/Differentiation/ArrayDifferentiation.swift

@@ -164,6 +164,14 @@ extension Array: Differentiable where Element: Differentiable {
     view.move(along: direction)
     self = view.base
   }
+
+  /// A closure that produces a `TangentVector` of zeros with the same
+  /// `count` as `self`.
+  public var zeroTangentVectorInitializer: () -> TangentVector {
+    { [count = self.count] in
+      TangentVector(.init(repeating: .zero, count: count))
+    }
+  }
 }
 
 //===----------------------------------------------------------------------===//

stdlib/public/Differentiation/CMakeLists.txt

@@ -14,6 +14,7 @@ add_swift_target_library(swift_Differentiation ${SWIFT_STDLIB_LIBRARY_BUILD_TYPE
   Differentiable.swift
   DifferentialOperators.swift
   DifferentiationUtilities.swift
+  AnyDerivative.swift
   ArrayDifferentiation.swift
 
   GYB_SOURCES

stdlib/public/Differentiation/Differentiable.swift

@@ -36,6 +36,40 @@ public protocol Differentiable {
   /// equivalent to exponential map, which moves `self` on the geodesic surface
   /// along the given tangent vector.
   mutating func move(along direction: TangentVector)
+
+  /// A closure that produces a zero tangent vector, capturing minimal
+  /// necessary information from `self`.
+  ///
+  /// `move(along: zeroTangentVectorInitializer())` should not modify
+  /// `self`.
+  ///
+  /// In some cases, the zero tangent vector of `self` is equal to
+  /// `TangentVector.zero`. In other cases, the zero tangent vector depends on
+  /// information in `self`, such as shape for an n-dimensional array type.
+  /// For differentiable programming, it is more memory-efficient to define a
+  /// custom `zeroTangentVectorInitializer` property which returns a closure
+  /// that captures and uses only the necessary information to create a zero
+  /// tangent vector. For example:
+  ///
+  ///     struct Vector {
+  ///         var scalars: [Float]
+  ///         var count: Int { scalars.count }
+  ///         init(scalars: [Float]) { ... }
+  ///         init(repeating repeatedElement: Float, count: Int) { ... }
+  ///     }
+  ///
+  ///     extension Vector: AdditiveArithmetic { ... }
+  ///
+  ///     extension Vector: Differentiable {
+  ///         typealias TangentVector = Vector
+  ///
+  ///         @noDerivative
+  ///         var zeroTangentVectorInitializer: () -> TangentVector {
+  ///             let count = self.count
+  ///             return { TangentVector(repeating: 0, count: count) }
+  ///         }
+  ///     }
+  var zeroTangentVectorInitializer: () -> TangentVector { get }
 }
 
 public extension Differentiable where TangentVector == Self {
@@ -44,3 +78,24 @@ public extension Differentiable where TangentVector == Self {
     self += direction
   }
 }
+
+public extension Differentiable {
+  // This is a temporary solution enabling the addition of
+  // `zeroTangentVectorInitializer` without implementing derived conformances.
+  // This property will produce incorrect results when tangent vectors depend
+  // on instance-specific information from `self`.
+  // TODO: Implement derived conformances and remove this default
+  // implementation.
+  @available(*, deprecated, message: """
+      `zeroTangentVectorInitializer` derivation has not been implemented; this \
+      default implementation is not correct when tangent vectors depend on \
+      instance-specific information from `self` and should not be used
+      """)
+  var zeroTangentVectorInitializer: () -> TangentVector {
+    { TangentVector.zero }
+  }
+
+  /// A tangent vector initialized using `zeroTangentVectorInitializer`.
+  /// `move(along: zeroTangentVector)` should not modify `self`.
+  var zeroTangentVector: TangentVector { zeroTangentVectorInitializer() }
+}