[AutoDiff] Deprecate Differentiable.AllDifferentiableVariables.

dan-zheng · dan-zheng · commit e3cc71f23624 · 2019-08-06T19:47:38.000-07:00
Deprecate the `AllDifferentiableVariables` associated type and `var allDifferentiableVariables` property of the `Differentiable` protocol. `AllDifferentiableVariables` is not essential for differentiable programming and was added as a workaround to enable key-path-based machine learning optimizers: let parameters and gradients have same type (`AllDifferentiableVariables == TangentVector`) to enable joint key-path iteration. It is possible to implement key-path-based machine learning optimizers via other means (do key-path-based operations on `TangentVector`, then call `Differentiable.move(along:)` to perform update), so `AllDifferentiableVariables` is no longer necessary. Resolves TF-707.
diff --git a/include/swift/AST/KnownIdentifiers.def b/include/swift/AST/KnownIdentifiers.def
@@ -155,9 +155,7 @@ IDENTIFIER(by)
 IDENTIFIER(scale)
 IDENTIFIER(x)
 // Differentiable
-IDENTIFIER(AllDifferentiableVariables)
 IDENTIFIER(TangentVector)
-IDENTIFIER(allDifferentiableVariables)
 IDENTIFIER(move)
 
 // Kinds of layout constraints
diff --git a/lib/Sema/DerivedConformanceDifferentiable.cpp b/lib/Sema/DerivedConformanceDifferentiable.cpp
diff --git a/lib/Sema/DerivedConformances.cpp b/lib/Sema/DerivedConformances.cpp
@@ -285,11 +285,6 @@ ValueDecl *DerivedConformance::getDerivableRequirement(NominalTypeDecl *nominal,
     if (name.isSimpleName(ctx.Id_typeList))
       return getRequirement(KnownProtocolKind::TensorGroup);
 
-    // SWIFT_ENABLE_TENSORFLOW
-    // Differentiable.allDifferentiableVariables
-    if (name.isSimpleName(ctx.Id_allDifferentiableVariables))
-      return getRequirement(KnownProtocolKind::Differentiable);
-
     return nullptr;
   }
 
@@ -454,9 +449,7 @@ ValueDecl *DerivedConformance::getDerivableRequirement(NominalTypeDecl *nominal,
 
     // SWIFT_ENABLE_TENSORFLOW
     // Differentiable.TangentVector
-    // Differentiable.AllDifferentiableVariables
-    if (name.isSimpleName(ctx.Id_TangentVector) ||
-        name.isSimpleName(ctx.Id_AllDifferentiableVariables))
+    if (name.isSimpleName(ctx.Id_TangentVector))
       return getRequirement(KnownProtocolKind::Differentiable);
 
     // SWIFT_ENABLE_TENSORFLOW
diff --git a/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift b/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift
@@ -65,20 +65,12 @@ public struct Tracked<T> {
   }
   private var handle: Box
 
-  @differentiable(
-    vjp: _vjpInit
-    where T : Differentiable, T == T.AllDifferentiableVariables,
-          T == T.TangentVector
-  )
+  @differentiable(vjp: _vjpInit where T : Differentiable, T == T.TangentVector)
   public init(_ value: T) {
     self.handle = Box(value)
   }
 
-  @differentiable(
-    vjp: _vjpValue
-    where T : Differentiable, T == T.AllDifferentiableVariables,
-          T == T.TangentVector
-  )
+  @differentiable(vjp: _vjpValue where T : Differentiable, T == T.TangentVector)
   public var value: T {
     get { handle.value }
     set { handle.value = newValue }
@@ -174,17 +166,11 @@ extension Tracked : Strideable where T : Strideable, T.Stride == T.Stride.Magnit
 }
 
 // For now, `T` must be restricted to trivial types (like `Float` or `Tensor`).
-extension Tracked : Differentiable
-  where T : Differentiable, T == T.AllDifferentiableVariables,
-        T == T.TangentVector
-{
-  public typealias AllDifferentiableVariables = Tracked<T.AllDifferentiableVariables>
+extension Tracked : Differentiable where T : Differentiable, T == T.TangentVector {
   public typealias TangentVector = Tracked<T.TangentVector>
 }
 
-extension Tracked where T : Differentiable, T == T.AllDifferentiableVariables,
-                        T == T.TangentVector
-{
+extension Tracked where T : Differentiable, T == T.TangentVector {
   @usableFromInline
   internal static func _vjpInit(_ value: T)
       -> (value: Self, pullback: (Self.TangentVector) -> (T.TangentVector)) {
@@ -197,9 +183,7 @@ extension Tracked where T : Differentiable, T == T.AllDifferentiableVariables,
   }
 }
 
-extension Tracked where T : Differentiable, T == T.AllDifferentiableVariables,
-                        T == T.TangentVector
-{
+extension Tracked where T : Differentiable, T == T.TangentVector {
   @usableFromInline
   @differentiating(+)
   internal static func _vjpAdd(lhs: Self, rhs: Self)
@@ -216,7 +200,7 @@ extension Tracked where T : Differentiable, T == T.AllDifferentiableVariables,
 }
 
 extension Tracked where T : Differentiable & SignedNumeric, T == T.Magnitude,
-                        T == T.AllDifferentiableVariables, T == T.TangentVector {
+                        T == T.TangentVector {
   @usableFromInline
   @differentiating(*)
   internal static func _vjpMultiply(lhs: Self, rhs: Self)
@@ -225,8 +209,7 @@ extension Tracked where T : Differentiable & SignedNumeric, T == T.Magnitude,
   }
 }
 
-extension Tracked where T : Differentiable & FloatingPoint,
-                        T == T.AllDifferentiableVariables, T == T.TangentVector {
+extension Tracked where T : Differentiable & FloatingPoint, T == T.TangentVector {
   @usableFromInline
   @differentiating(/)
   internal static func _vjpDivide(lhs: Self, rhs: Self)
diff --git a/stdlib/public/core/Array.swift b/stdlib/public/core/Array.swift
@@ -1952,26 +1952,6 @@ extension Array where Element : Differentiable {
 
     public typealias TangentVector =
       Array<Element.TangentVector>.DifferentiableView
-    public typealias AllDifferentiableVariables =
-      Array<Element.AllDifferentiableVariables>.DifferentiableView
-
-    public var allDifferentiableVariables: AllDifferentiableVariables {
-      get {
-        return AllDifferentiableVariables(
-          base.map { $0.allDifferentiableVariables })
-      }
-      set {
-        precondition(
-          base.count == newValue.base.count,
-          "cannot set Array.DifferentiableView.AllDifferentiableVariables " +
-            "with count \(base.count) to " +
-            "Array.DifferentiableView.AllDifferentiableVariables with " +
-            "different count \(newValue.base.count)")
-        for i in base.indices {
-          base[i].allDifferentiableVariables = newValue.base[i]
-        }
-      }
-    }
 
     public mutating func move(along direction: TangentVector) {
       precondition(
@@ -2066,27 +2046,13 @@ extension Array.DifferentiableView : AdditiveArithmetic
 /// Makes `Array` differentiable as the product manifold of `Element`
 /// multiplied with itself `count` times.
 extension Array : Differentiable where Element : Differentiable {
-  // In an ideal world, `TangentVector`, `TangentVector`, and
-  // `AllDifferentiableVariables` would all be `Array`s. Unfortunately, we
-  // can't conform `Array` to `AdditiveArithmetic` for `TangentVector` and
-  // `TangentVector`, because `Array` already has a static `+` method with
-  // different semantics from `AdditiveArithmetic` `+`. So we use
+  // In an ideal world, `TangentVector` would be `[Element.TangentVector]`.
+  // Unfortunately, we cannot conform `Array` to `AdditiveArithmetic` for
+  // `TangentVector` because `Array` already has a static `+` method with
+  // different semantics from `AdditiveArithmetic.+`. So we use
   // `Array.DifferentiableView` for all these associated types.
   public typealias TangentVector =
     Array<Element.TangentVector>.DifferentiableView
-  public typealias AllDifferentiableVariables =
-    Array<Element.AllDifferentiableVariables>.DifferentiableView
-
-  public var allDifferentiableVariables: AllDifferentiableVariables {
-    get {
-      return DifferentiableView(self).allDifferentiableVariables
-    }
-    set {
-      var view = DifferentiableView(self)
-      view.allDifferentiableVariables = newValue
-      self = view.base
-    }
-  }
 
   public mutating func move(along direction: TangentVector) {
     var view = DifferentiableView(self)
diff --git a/stdlib/public/core/AutoDiff.swift b/stdlib/public/core/AutoDiff.swift
@@ -153,26 +153,24 @@ public extension VectorProtocol where VectorSpaceScalar : SignedNumeric {
 /// tangent spaces are finite-dimensional.
 public protocol Differentiable {
   associatedtype TangentVector: Differentiable & AdditiveArithmetic
-    where TangentVector.TangentVector == TangentVector,
-          AllDifferentiableVariables.AllDifferentiableVariables ==
-            AllDifferentiableVariables,
-          AllDifferentiableVariables.TangentVector == TangentVector
-  /// The type of all differentiable variables in this type.
-  associatedtype AllDifferentiableVariables : Differentiable
-
-  /// All differentiable variables of this value.
-  var allDifferentiableVariables: AllDifferentiableVariables { get set }
+    where TangentVector.TangentVector == TangentVector
 
   /// Moves `self` along the value space towards the given tangent vector. In
   /// Riemannian geometry (mathematics), this represents an exponential map.
   mutating func move(along direction: TangentVector)
 
+  @available(*, deprecated,
+             message: "'AllDifferentiableVariables' is now equal to 'Self' and will be removed")
+  typealias AllDifferentiableVariables = Self
+
   @available(*, deprecated,
              message: "'CotangentVector' is now equal to 'TangentVector' and will be removed")
   typealias CotangentVector = TangentVector
 }
 
-public extension Differentiable where AllDifferentiableVariables == Self {
+public extension Differentiable {
+  @available(*, deprecated,
+             message: "'allDifferentiableVariables' is now equal to 'self' and will be removed")
   var allDifferentiableVariables: AllDifferentiableVariables {
     get { return self }
     set { self = newValue }
@@ -731,8 +729,7 @@ internal protocol _AnyDerivativeBox {
 
   /// 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,
-          U.AllDifferentiableVariables == U
+    where U : Differentiable, U.TangentVector == U
 }
 
 extension _AnyDerivativeBox {
@@ -754,8 +751,7 @@ internal func _derivativeTypeMismatch(
 }
 
 internal struct _ConcreteDerivativeBox<T> : _AnyDerivativeBox
-  where T : Differentiable, T.TangentVector == T,
-        T.AllDifferentiableVariables == T
+  where T : Differentiable, T.TangentVector == T
 {
   /// The underlying base value.
   var _base: T
@@ -770,8 +766,7 @@ internal struct _ConcreteDerivativeBox<T> : _AnyDerivativeBox
   }
 
   func _unboxed<U>(to type: U.Type) -> U?
-    where U : Differentiable, U.TangentVector == U,
-          U.AllDifferentiableVariables == U
+    where U : Differentiable, U.TangentVector == U
   {
     return (self as? _ConcreteDerivativeBox<U>)?._base
   }
@@ -861,24 +856,19 @@ public struct AnyDerivative : Differentiable & AdditiveArithmetic {
 
   /// Creates a type-erased derivative from the given derivative.
   @differentiable(vjp: _vjpInit(_:))
-  public init<T>(_ base: T)
-    where T : Differentiable, T.TangentVector == T,
-          T.AllDifferentiableVariables == T
-  {
+  public init<T>(_ base: T) where T : Differentiable, T.TangentVector == T {
     self._box = _ConcreteDerivativeBox<T>(base)
   }
 
   @usableFromInline internal static func _vjpInit<T>(
     _ base: T
   ) -> (AnyDerivative, (AnyDerivative) -> T.TangentVector)
-    where T : Differentiable, T.TangentVector == T,
-          T.AllDifferentiableVariables == T
+    where T : Differentiable, T.TangentVector == T
   {
     return (AnyDerivative(base), { v in v.base as! T.TangentVector })
   }
 
   public typealias TangentVector = AnyDerivative
-  public typealias AllDifferentiableVariables = AnyDerivative
 
   // `Equatable` requirements (implied by `AdditiveArithmetic`).
   public static func == (lhs: AnyDerivative, rhs: AnyDerivative) -> Bool {
@@ -929,10 +919,6 @@ public struct AnyDerivative : Differentiable & AdditiveArithmetic {
   }
 
   // `Differentiable` requirements.
-  public var allDifferentiableVariables: AllDifferentiableVariables {
-    get { return AnyDerivative(_box: _box._allDifferentiableVariables) }
-    // set { _box._allDifferentiableVariables = newValue._box }
-  }
   public mutating func move(along direction: TangentVector) {
     if _box._isOpaqueZero() {
       _box = direction._box
diff --git a/stdlib/public/core/FloatingPointTypes.swift.gyb b/stdlib/public/core/FloatingPointTypes.swift.gyb
@@ -1906,7 +1906,6 @@ extension ${Self} : VectorProtocol {
 
 extension ${Self} : Differentiable {
   public typealias TangentVector = ${Self}
-  public typealias AllDifferentiableVariables = ${Self}
 
   public mutating func move(along direction: TangentVector) {
     self += direction
diff --git a/stdlib/public/core/SIMDVectorTypes.swift.gyb b/stdlib/public/core/SIMDVectorTypes.swift.gyb
@@ -196,10 +196,6 @@ extension SIMD${n} : Differentiable
   where Scalar : Differentiable & BinaryFloatingPoint,
         Scalar.TangentVector : BinaryFloatingPoint {
   public typealias TangentVector = SIMD${n}
-  public typealias AllDifferentiableVariables = SIMD${n}
-  public func tangentVector(from cotangent: TangentVector) -> TangentVector {
-    return cotangent
-  }
 }
 
 extension SIMD${n}
diff --git a/test/Sema/class_differentiable.swift b/test/Sema/class_differentiable.swift
diff --git a/test/Sema/struct_differentiable.swift b/test/Sema/struct_differentiable.swift
