[AutoDiff] Add 'zeroTangentVector' property to 'Differentiable' protocol.

Zero tangent vector is necessary for optimizations on models with an array of parameters, especially for optimizers that iterates over parameters using key paths. The [current implementation](https://github.com/tensorflow/swift-apis/blob/master/Sources/TensorFlow/Optimizers/MomentumBased.swift) of some key-path-based optimizers is wrong in that it won't work with models that contain an array of parameters (tangent vectors like `infinityNorm` are initialized as `.zero`).

An earlier version of these optimizer using the deprecated `AllDifferentiableVariables` property would give the correct results, but would be heavyweight and inefficient because they'd need to 1. add a constraint `TangentVector == AllDifferentiableVariables` to optimizers, and 2. make a copy of all parameters and resetting them to `.zero`. Since we are deprecating `AllDifferentiableVariables`, this is not the right direction.

This problem also means that our `Differentiable` abstraction needs to provide a general mechanism of obtaining a zero tangent vector at a certain instance. Hence we add a `zeroTangentVector` property to the `Differentiable` protocol.

Zero tangent vectors do not have a canonical mathematical definition, but makes sense for `Differentiable` in the standard library because Swift does not have dependent types and thus cannot have a `TangentVector` that depends on a point on a differentiable manifold. Manopt also has an API, `M.zerovec(x)`, that creates a zero tangent vector at a point (see their API doc [here](https://www.manopt.org/tutorial.html).

Adding `zeroTangentVector` will make it possible to deprecate `AllDifferentiableVariables` completely, because currently some fast.ai notebooks depend on initializing parameter gradients using `AllDifferentiableVariables`.

The new `Differentiable` protocol looks like the following. The [design overview](http://bit.ly/swift-autodiff) has been updated to reflect this change.

```swift
protocol Differentiable {
    /// A type representing the differentiable value’s derivatives.
    /// Mathematically, this is equivalent to the tangent bundle of the
    /// differentiable manifold represented by the differentiable type.
    associatedtype TangentVector: Differentiable & AdditiveArithmetic

    /// Moves `self` along the given direction. In Riemannian geometry,
    /// this is equivalent to exponential map, which moves `self` on the
    /// geodesic surface along the given tangent vector.
    mutating func move(along direction: TangentVector)

    /// A tangent vector such that `move(along: zeroTangentVector)`
    /// will not modify `self`.
    /// - Note: `zeroTangentVector` can be `TangentVector.zero` in most cases,
    ///   but types whose tangent vectors depend on instance properties of
    ///   `self` need to provide a different implementation. For example, an
    ///   array’s zero tangent vector depends on the array’s `count`.
    var zeroTangentVector: TangentVector { get }
}
```

Author: rxwei

include/swift/AST/KnownIdentifiers.def

@@ -159,6 +159,7 @@ IDENTIFIER(AllDifferentiableVariables)
 IDENTIFIER(TangentVector)
 IDENTIFIER(allDifferentiableVariables)
 IDENTIFIER(move)
+IDENTIFIER(zeroTangentVector)
 
 // Kinds of layout constraints
 IDENTIFIER_WITH_NAME(UnknownLayout, "_UnknownLayout")

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -380,6 +380,105 @@ static ValueDecl *getUnderlyingAllDiffableVariables(DeclContext *DC,
   return allDiffableVarsDecl;
 }
 
+// Return the underlying `zeroTangentVector` of a VarDecl `x`. If `x` conforms
+// to `Differentiable`, return `zeroTangentVector`. Otherwise, return
+// `x`.
+static ValueDecl *getUnderlyingZeroTangentVector(DeclContext *DC,
+                                                 VarDecl *varDecl) {
+  auto *module = DC->getParentModule();
+  auto &C = module->getASTContext();
+  auto *diffableProto = C.getProtocol(KnownProtocolKind::Differentiable);
+  auto *zeroTangentVectorReq =
+      getProtocolRequirement(diffableProto, C.Id_zeroTangentVector);
+  if (!varDecl->hasInterfaceType())
+    C.getLazyResolver()->resolveDeclSignature(varDecl);
+  auto varType = DC->mapTypeIntoContext(varDecl->getValueInterfaceType());
+  auto confRef = module->lookupConformance(varType, diffableProto);
+  if (!confRef)
+    return varDecl;
+  // Use protocol requirement as a default for abstract conformances.
+  // If conformance is concrete, get concrete witness declaration instead.
+  ValueDecl *zeroTangentVectorDecl = zeroTangentVectorReq;
+  if (confRef->isConcrete())
+    zeroTangentVectorDecl = confRef->getConcrete()->getWitnessDecl(
+        zeroTangentVectorReq);
+  return zeroTangentVectorDecl;
+}
+
+// Get the effective memberwise initializer of the given nominal type, or create
+// it if it does not exist.
+static ConstructorDecl *getOrCreateEffectiveMemberwiseInitializer(
+    TypeChecker &TC, NominalTypeDecl *nominal) {
+  auto &C = nominal->getASTContext();
+  if (auto *initDecl = nominal->getEffectiveMemberwiseInitializer())
+    return initDecl;
+  auto *initDecl = createImplicitConstructor(
+      TC, nominal, ImplicitConstructorKind::Memberwise);
+  nominal->addMember(initDecl);
+  C.addSynthesizedDecl(initDecl);
+  return initDecl;
+}
+
+// Synthesize getter body for `zeroTangentVector` computed property.
+static std::pair<BraceStmt *, bool>
+derivedBody_zeroTangentVectorGetter(AbstractFunctionDecl *getterDecl, void *) {
+  auto *parentDC = getterDecl->getParent();
+  auto *nominal = parentDC->getSelfNominalTypeDecl();
+  auto &C = nominal->getASTContext();
+
+  auto *tangentVectorStruct =
+      getAssociatedStructDecl(parentDC, C.Id_TangentVector);
+  auto *tangentVectorInitDecl =
+      tangentVectorStruct->getEffectiveMemberwiseInitializer();
+  assert(tangentVectorInitDecl &&
+         "'TangentVector' memberwise initializer not found");
+  auto *selfDecl = getterDecl->getImplicitSelfDecl();
+  auto *selfDRE =
+      new (C) DeclRefExpr(selfDecl, DeclNameLoc(), /*Implicit*/ true);
+
+  auto *initDRE = new (C) DeclRefExpr(tangentVectorInitDecl, DeclNameLoc(),
+                                      /*Implicit*/ true);
+  initDRE->setFunctionRefKind(FunctionRefKind::SingleApply);
+
+  auto tangentVectorType = parentDC->mapTypeIntoContext(
+      tangentVectorStruct->getDeclaredInterfaceType());
+  Expr *baseExpr = TypeExpr::createImplicit(tangentVectorType, C);
+  auto *initExpr = new (C) ConstructorRefCallExpr(initDRE, baseExpr);
+  initExpr->setThrows(false);
+  initExpr->setImplicit();
+
+  SmallVector<Expr *, 2> members;
+  SmallVector<Identifier, 2> memberNames;
+
+  llvm::DenseMap<Identifier, VarDecl *> diffPropertyMap;
+  SmallVector<VarDecl *, 8> diffProperties;
+  getStoredPropertiesForDifferentiation(nominal, parentDC, diffProperties);
+  for (auto *member : diffProperties)
+    diffPropertyMap[member->getName()] = member;
+
+  for (auto initParam : *tangentVectorInitDecl->getParameters()) {
+    auto member = diffPropertyMap[initParam->getName()];
+    member->setInterfaceType(member->getValueInterfaceType());
+    Expr *memberExpr = new (C) MemberRefExpr(selfDRE, SourceLoc(), member,
+                                             DeclNameLoc(), /*Implicit*/ true);
+    auto *memberZeroTangentVector =
+        getUnderlyingZeroTangentVector(parentDC, member);
+    auto *memberZeroTangentExpr =
+        new (C) MemberRefExpr(memberExpr, SourceLoc(), memberZeroTangentVector,
+                              DeclNameLoc(), /*Implicit*/ true);
+    members.push_back(memberZeroTangentExpr);
+    memberNames.push_back(member->getName());
+  }
+  Expr *callExpr = CallExpr::createImplicit(C, initExpr, members, memberNames);
+
+  ASTNode returnStmt =
+      new (C) ReturnStmt(SourceLoc(), callExpr, /*Implicit*/ true);
+  auto *braceStmt =
+      BraceStmt::create(C, SourceLoc(), returnStmt, SourceLoc(),
+                        /*Implicit*/ true);
+  return std::make_pair(braceStmt, false);
+}
+
 // Synthesize getter body for `allDifferentiableVariables` computed property.
 static std::pair<BraceStmt *, bool>
 derivedBody_allDifferentiableVariablesGetter(AbstractFunctionDecl *getterDecl,
@@ -537,6 +636,41 @@ deriveDifferentiable_allDifferentiableVariables(DerivedConformance &derived) {
   return allDiffableVarsDecl;
 }
 
+// Synthesize `zeroTangentVector` computed property declaration.
+static ValueDecl *
+deriveDifferentiable_zeroTangentVector(DerivedConformance &derived) {
+  auto *parentDC = derived.getConformanceContext();
+  auto &TC = derived.TC;
+  auto &C = TC.Context;
+
+  // Get `TangentVector` struct.
+  auto *tangentVectorStruct =
+      getAssociatedStructDecl(parentDC, C.Id_TangentVector);
+  // Make sure a memberwise initializer exists because the body synthesizer
+  // needs it.
+  getOrCreateEffectiveMemberwiseInitializer(TC, tangentVectorStruct);
+
+  auto returnInterfaceTy = tangentVectorStruct->getDeclaredInterfaceType();
+  auto returnTy = parentDC->mapTypeIntoContext(returnInterfaceTy);
+
+  VarDecl *zeroTangentVectorDecl;
+  PatternBindingDecl *pbDecl;
+  std::tie(zeroTangentVectorDecl, pbDecl) = derived.declareDerivedProperty(
+      C.Id_zeroTangentVector, returnInterfaceTy, returnTy,
+      /*isStatic*/ false, /*isFinal*/ true);
+
+  auto *getterDecl = derived.addGetterToReadOnlyDerivedProperty(
+      zeroTangentVectorDecl, returnTy);
+  getterDecl->setBodySynthesizer(&derivedBody_zeroTangentVectorGetter);
+  derived.addMembersToConformanceContext(
+      {getterDecl, zeroTangentVectorDecl, pbDecl});
+
+  addExpectedOpaqueAccessorsToStorage(zeroTangentVectorDecl, C);
+  triggerAccessorSynthesis(TC, zeroTangentVectorDecl);
+
+  return zeroTangentVectorDecl;
+}
+
 // Return associated `TangentVector` or `AllDifferentiableVariables` struct for
 // a nominal type, if it exists.
 // If not, synthesize the struct. Also return a Boolean value that indicates
@@ -1035,6 +1169,8 @@ ValueDecl *DerivedConformance::deriveDifferentiable(ValueDecl *requirement) {
     return deriveDifferentiable_move(*this);
   if (requirement->getBaseName() == TC.Context.Id_allDifferentiableVariables)
     return deriveDifferentiable_allDifferentiableVariables(*this);
+  if (requirement->getBaseName() == TC.Context.Id_zeroTangentVector)
+    return deriveDifferentiable_zeroTangentVector(*this);
   TC.diagnose(requirement->getLoc(), diag::broken_differentiable_requirement);
   return nullptr;
 }

lib/Sema/DerivedConformances.cpp

@@ -290,6 +290,11 @@ ValueDecl *DerivedConformance::getDerivableRequirement(NominalTypeDecl *nominal,
     if (name.isSimpleName(ctx.Id_allDifferentiableVariables))
       return getRequirement(KnownProtocolKind::Differentiable);
 
+    // SWIFT_ENABLE_TENSORFLOW
+    // Differentiable.zeroTangentVector
+    if (name.isSimpleName(ctx.Id_zeroTangentVector))
+      return getRequirement(KnownProtocolKind::Differentiable);
+
     return nullptr;
   }
 

stdlib/public/core/Array.swift

@@ -1982,6 +1982,11 @@ extension Array where Element : Differentiable {
         base[i].move(along: direction.base[i])
       }
     }
+
+    public var zeroTangentVector: TangentVector {
+      TangentVector(Array<Element.TangentVector>(repeating: .zero,
+                                                 count: base.count))
+    }
   }
 }
 
@@ -2093,6 +2098,10 @@ extension Array : Differentiable where Element : Differentiable {
     view.move(along: direction)
     self = view.base
   }
+
+  public var zeroTangentVector: TangentVector {
+    TangentVector(Array<Element.TangentVector>(repeating: .zero, count: count))
+  }
 }
 
 extension Array where Element : Differentiable {

stdlib/public/core/AutoDiff.swift

@@ -152,6 +152,9 @@ public extension VectorProtocol where VectorSpaceScalar : SignedNumeric {
 /// A type that mathematically represents a differentiable manifold whose
 /// tangent spaces are finite-dimensional.
 public protocol Differentiable {
+  /// A type representing a differentiable value’s derivatives.
+  /// Mathematically, this is equivalent to the tangent bundle of the
+  /// differentiable manifold represented by the differentiable type.
   associatedtype TangentVector: Differentiable & AdditiveArithmetic
     where TangentVector.TangentVector == TangentVector,
           AllDifferentiableVariables.AllDifferentiableVariables ==
@@ -163,10 +166,20 @@ public protocol Differentiable {
   /// All differentiable variables of this value.
   var allDifferentiableVariables: AllDifferentiableVariables { get set }
 
-  /// Moves `self` along the value space towards the given tangent vector. In
-  /// Riemannian geometry (mathematics), this represents an exponential map.
+  /// Moves `self` along the given direction. In Riemannian geometry,
+  /// this is equivalent to exponential map, which moves `self` on the
+  /// geodesic surface along the given tangent vector.
   mutating func move(along direction: TangentVector)
 
+  /// A tangent vector such that `move(along: zeroTangentVector)` will not
+  /// modify `self`.
+  ///
+  /// - Note: `zeroTangentVector` can be `TangentVector.zero` in most cases,
+  ///   but types whose tangent vectors depend on instance properties of
+  ///   `self` need to provide a different implementation. For example, an
+  ///   array’s zero tangent vector depends on the array’s `count`.
+  var zeroTangentVector: TangentVector { get }
+
   @available(*, deprecated,
              message: "'CotangentVector' is now equal to 'TangentVector' and will be removed")
   typealias CotangentVector = TangentVector
@@ -183,6 +196,9 @@ public extension Differentiable where TangentVector == Self {
   mutating func move(along direction: TangentVector) {
     self += direction
   }
+  var zeroTangentVector: TangentVector {
+    return .zero
+  }
 }
 
 /// Returns `x` like an identity function. When used in a context where `x` is
@@ -725,6 +741,7 @@ internal protocol _AnyDerivativeBox {
   // `Differentiable` requirements.
   var _allDifferentiableVariables: _AnyDerivativeBox { get }
   mutating func _move(along direction: _AnyDerivativeBox)
+  var _zeroTangentVector: _AnyDerivativeBox { get }
 
   /// The underlying base value, type-erased to `Any`.
   var _typeErasedBase: Any { get }
@@ -840,6 +857,10 @@ internal struct _ConcreteDerivativeBox<T> : _AnyDerivativeBox
     }
     _base.move(along: directionBase)
   }
+
+  var _zeroTangentVector: _AnyDerivativeBox {
+    return _ConcreteDerivativeBox(_base.zeroTangentVector)
+  }
 }
 
 /// A type-erased derivative value.
@@ -940,6 +961,9 @@ public struct AnyDerivative : Differentiable & AdditiveArithmetic {
     }
     _box._move(along: direction._box)
   }
+  public var zeroTangentVector: TangentVector {
+    AnyDerivative(_box: _box._zeroTangentVector)
+  }
 }
 
 //===----------------------------------------------------------------------===//

stdlib/public/core/FloatingPointTypes.swift.gyb

@@ -1911,6 +1911,8 @@ extension ${Self} : Differentiable {
   public mutating func move(along direction: TangentVector) {
     self += direction
   }
+
+  public var zeroTangentVector: TangentVector { .zero }
 }
 
 //===----------------------------------------------------------------------===//

stdlib/public/core/SIMDVectorTypes.swift.gyb

@@ -200,6 +200,7 @@ extension SIMD${n} : Differentiable
   public func tangentVector(from cotangent: TangentVector) -> TangentVector {
     return cotangent
   }
+  public var zeroTangentVector: TangentVector { .zero }
 }
 
 extension SIMD${n}

test/Sema/struct_differentiable.swift

@@ -320,6 +320,7 @@ struct VectorSpaceTypeAlias : AdditiveArithmetic, Differentiable {
   var w: Float
   var b: Float
   typealias TangentVector = Simple
+  var zeroTangentVector: TangentVector { .zero }
 }
 // expected-error @+2 {{type 'VectorSpaceCustomStruct' does not conform to protocol 'Differentiable'}}
 // expected-note @+1 {{do you want to add protocol stubs?}}
@@ -331,6 +332,7 @@ struct VectorSpaceCustomStruct : AdditiveArithmetic, Differentiable {
     var b: Float.TangentVector
     typealias TangentVector = VectorSpaceCustomStruct.TangentVector
   }
+  var zeroTangentVector: TangentVector { .zero }
 }
 
 struct StaticNoDerivative : Differentiable {
@@ -371,14 +373,10 @@ extension NoMemberwiseInitializerExtended: Differentiable
 
 // Test derived conformances in disallowed contexts.
 
-// expected-error @+4 {{type 'OtherFileNonconforming' does not conform to protocol 'Differentiable'}}
-// expected-error @+3 {{implementation of 'Differentiable' cannot be automatically synthesized in an extension in a different file to the type}}
-// expected-note @+2 {{do you want to add protocol stubs?}}
-// expected-note @+1 {{do you want to add protocol stubs?}}
+// expected-error @+2 {{type 'OtherFileNonconforming' does not conform to protocol 'Differentiable'}}
+// expected-error @+1 {{implementation of 'Differentiable' cannot be automatically synthesized in an extension in a different file to the type}}
 extension OtherFileNonconforming : Differentiable {}
 
-// expected-error @+4 {{type 'GenericOtherFileNonconforming<T>' does not conform to protocol 'Differentiable'}}
-// expected-error @+3 {{implementation of 'Differentiable' cannot be automatically synthesized in an extension in a different file to the type}}
-// expected-note @+2 {{do you want to add protocol stubs?}}
-// expected-note @+1 {{do you want to add protocol stubs?}}
+// expected-error @+2 {{type 'GenericOtherFileNonconforming<T>' does not conform to protocol 'Differentiable'}}
+// expected-error @+1 {{implementation of 'Differentiable' cannot be automatically synthesized in an extension in a different file to the type}}
 extension GenericOtherFileNonconforming : Differentiable {}