Merge pull request #31781 from dan-zheng/array-append-derivative

[AutoDiff] Fix `Array.append(_:)` pullback.

Author: dan-zheng

stdlib/public/Differentiation/ArrayDifferentiation.swift

@@ -184,36 +184,31 @@ extension Array where Element: Differentiable {
   func _vjpSubscript(index: Int) -> (
     value: Element, pullback: (Element.TangentVector) -> TangentVector
   ) {
-    func pullback(_ gradientIn: Element.TangentVector) -> TangentVector {
-      var gradientOut = [Element.TangentVector](
+    func pullback(_ v: Element.TangentVector) -> TangentVector {
+      var dSelf = [Element.TangentVector](
         repeating: .zero,
         count: count)
-      gradientOut[index] = gradientIn
-      return TangentVector(gradientOut)
+      dSelf[index] = v
+      return TangentVector(dSelf)
     }
     return (self[index], pullback)
   }
 
   @usableFromInline
   @derivative(of: +)
-  static func _vjpConcatenate(_ lhs: [Element], _ rhs: [Element]) -> (
-    value: [Element],
+  static func _vjpConcatenate(_ lhs: Self, _ rhs: Self) -> (
+    value: Self,
     pullback: (TangentVector) -> (TangentVector, TangentVector)
   ) {
-    func pullback(_ gradientIn: TangentVector) -> (TangentVector, TangentVector)
-    {
+    func pullback(_ v: TangentVector) -> (TangentVector, TangentVector) {
       precondition(
-        gradientIn.base.count == lhs.count + rhs.count,
+        v.base.count == lhs.count + rhs.count,
         "+ should receive gradient with count equal to sum of operand "
-          + "counts, but counts are: gradient \(gradientIn.base.count), "
+          + "counts, but counts are: gradient \(v.base.count), "
           + "lhs \(lhs.count), rhs \(rhs.count)")
       return (
-        TangentVector(
-          [Element.TangentVector](
-            gradientIn.base[0..<lhs.count])),
-        TangentVector(
-          [Element.TangentVector](
-            gradientIn.base[lhs.count...]))
+        TangentVector([Element.TangentVector](v.base[0..<lhs.count])),
+        TangentVector([Element.TangentVector](v.base[lhs.count...]))
       )
     }
     return (lhs + rhs, pullback)
@@ -227,8 +222,11 @@ extension Array where Element: Differentiable {
     value: Void, pullback: (inout TangentVector) -> Element.TangentVector
   ) {
     let appendedElementIndex = count
-    defer { append(element) }
-    return ((), { dself in dself.base[appendedElementIndex] })
+    append(element)
+    return ((), { v in
+      defer { v.base.removeLast() }
+      return v.base[appendedElementIndex]
+    })
   }
 
   @usableFromInline
@@ -288,6 +286,37 @@ extension Array where Element: Differentiable {
 // Differentiable higher order functions for collections
 //===----------------------------------------------------------------------===//
 
+extension Array where Element: Differentiable {
+  @inlinable
+  @differentiable(wrt: self)
+  public func differentiableMap<Result: Differentiable>(
+    _ body: @differentiable (Element) -> Result
+  ) -> [Result] {
+    map(body)
+  }
+
+  @inlinable
+  @derivative(of: differentiableMap)
+  internal func _vjpDifferentiableMap<Result: Differentiable>(
+    _ body: @differentiable (Element) -> Result
+  ) -> (
+    value: [Result],
+    pullback: (Array<Result>.TangentVector) -> Array.TangentVector
+  ) {
+    var values: [Result] = []
+    var pullbacks: [(Result.TangentVector) -> Element.TangentVector] = []
+    for x in self {
+      let (y, pb) = valueWithPullback(at: x, in: body)
+      values.append(y)
+      pullbacks.append(pb)
+    }
+    func pullback(_ tans: Array<Result>.TangentVector) -> Array.TangentVector {
+      .init(zip(tans.base, pullbacks).map { tan, pb in pb(tan) })
+    }
+    return (value: values, pullback: pullback)
+  }
+}
+
 extension Array where Element: Differentiable {
   @inlinable
   @differentiable(wrt: (self, initialResult))
@@ -336,34 +365,3 @@ extension Array where Element: Differentiable {
     )
   }
 }
-
-extension Array where Element: Differentiable {
-  @inlinable
-  @differentiable(wrt: self)
-  public func differentiableMap<Result: Differentiable>(
-    _ body: @differentiable (Element) -> Result
-  ) -> [Result] {
-    map(body)
-  }
-
-  @inlinable
-  @derivative(of: differentiableMap)
-  internal func _vjpDifferentiableMap<Result: Differentiable>(
-    _ body: @differentiable (Element) -> Result
-  ) -> (
-    value: [Result],
-    pullback: (Array<Result>.TangentVector) -> Array.TangentVector
-  ) {
-    var values: [Result] = []
-    var pullbacks: [(Result.TangentVector) -> Element.TangentVector] = []
-    for x in self {
-      let (y, pb) = valueWithPullback(at: x, in: body)
-      values.append(y)
-      pullbacks.append(pb)
-    }
-    func pullback(_ tans: Array<Result>.TangentVector) -> Array.TangentVector {
-      .init(zip(tans.base, pullbacks).map { tan, pb in pb(tan) })
-    }
-    return (value: values, pullback: pullback)
-  }
-}

test/AutoDiff/stdlib/array.swift renamed to test/AutoDiff/validation-test/array.swift

@@ -372,15 +372,28 @@ ArrayAutoDiffTests.test("Array.+=") {
 }
 
 ArrayAutoDiffTests.test("Array.append") {
+  func appending(_ array: [Float], _ element: Float) -> [Float] {
+    var result = array
+    result.append(element)
+    return result
+  }
+  do {
+    let v = FloatArrayTan([1, 2, 3, 4])
+    expectEqual((.init([1, 2, 3]), 4),
+                pullback(at: [0, 0, 0], 0, in: appending)(v))
+  }
+
   func identity(_ array: [Float]) -> [Float] {
     var results: [Float] = []
     for i in withoutDerivative(at: array.indices) {
       results.append(array[i])
     }
     return results
   }
-  let v = FloatArrayTan([4, -5, 6])
-  expectEqual(v, pullback(at: [1, 2, 3], in: identity)(v))
+  do {
+    let v = FloatArrayTan([4, -5, 6])
+    expectEqual(v, pullback(at: [1, 2, 3], in: identity)(v))
+  }
 }
 
 ArrayAutoDiffTests.test("Array.init(repeating:count:)") {