[AutoDiff upstream] Add Differentiable.zeroTangentVector.

Add `Differentiable.zeroTangentVectorInitializer` protocol requirement and
`Differentiable.zeroTangentVector` default implementation.

Author: dan-zheng

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/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() }
+}

test/AutoDiff/stdlib/array.swift

@@ -401,13 +401,10 @@ ArrayAutoDiffTests.test("Array.DifferentiableView.base") {
     backprop(FloatArrayTan([1, 2, 3, 4])))
 }
 
-// TODO: Upstream `Differentiable.zeroTangentVector` and implementations.
-/*
 ArrayAutoDiffTests.test("Array.zeroTangentVector") {
   let count = 10
   let array: [Float] = Array((0..<count).map(Float.init))
   expectEqual(array.zeroTangentVector.base, Array(repeating: 0, count: count))
 }
-*/
 
 runAllTests()