[AutoDiff] [stdlib] Add differentiable map and reduce methods on 'Array'. (#26023)

Add variants of map and reduce that take a `@differentiable` closure and are themselves differentiable.
```swift
extension Array {
  @differentiable(wrt: self)
  func differentiableMap<Result: Differentiable>(
    _ body: @differentiable (Element) -> Result
  ) -> [Result]

  @differentiable(wrt: (self, initialResult))
  func differentiableReduce<Result: Differentiable>(
    _ initialResult: Result,
    _ nextPartialResult: @differentiable (Result, Element) -> Result
  ) -> Result
}
```

Also make `Array.DifferentiableView` conform to `ExpressibleByArrayLiteral` so that tests and user code are easier to write.

Author: rxwei

stdlib/public/core/Array.swift

@@ -1994,6 +1994,12 @@ extension Array.DifferentiableView : Equatable where Element : Equatable {
   }
 }
 
+extension Array.DifferentiableView : ExpressibleByArrayLiteral {
+  public init(arrayLiteral elements: Element...) {
+    self.init(elements)
+  }
+}
+
 extension Array.DifferentiableView : CustomStringConvertible {
   public var description: String {
     return base.description

stdlib/public/core/AutoDiff.swift

@@ -885,3 +885,76 @@ public struct AnyDerivative : Differentiable & AdditiveArithmetic {
     _box._move(along: direction._box)
   }
 }
+
+//===----------------------------------------------------------------------===//
+// Differentiable higher order functions for collections
+//===----------------------------------------------------------------------===//
+
+public extension Array where Element: Differentiable {
+  @differentiable(wrt: (self, initialResult), vjp: _vjpDifferentiableReduce)
+  func differentiableReduce<Result: Differentiable>(
+    _ initialResult: Result,
+    _ nextPartialResult: @differentiable (Result, Element) -> Result
+  ) -> Result {
+    reduce(initialResult, nextPartialResult)
+  }
+
+  @usableFromInline
+  internal func _vjpDifferentiableReduce<Result: Differentiable>(
+    _ initialResult: Result,
+    _ nextPartialResult: @differentiable (Result, Element) -> Result
+  ) -> (value: Result,
+        pullback: (Result.TangentVector)
+          -> (Array.TangentVector, Result.TangentVector)) {
+    var pullbacks:
+      [(Result.TangentVector) -> (Result.TangentVector, Element.TangentVector)]
+        = []
+    let count = self.count
+    pullbacks.reserveCapacity(count)
+    var result = initialResult
+    for element in self {
+      let (y, pb) =
+        Swift.valueWithPullback(at: result, element, in: nextPartialResult)
+      result = y
+      pullbacks.append(pb)
+    }
+    return (value: result, pullback: { tangent in
+      var resultTangent = tangent
+      var elementTangents = TangentVector([])
+      elementTangents.base.reserveCapacity(count)
+      for pullback in pullbacks.reversed() {
+        let (newResultTangent, elementTangent) = pullback(resultTangent)
+        resultTangent = newResultTangent
+        elementTangents.base.append(elementTangent)
+      }
+      return (TangentVector(elementTangents.base.reversed()), resultTangent)
+    })
+  }
+}
+
+public extension Array where Element: Differentiable {
+  @differentiable(wrt: self, vjp: _vjpDifferentiableMap)
+  func differentiableMap<Result: Differentiable>(
+    _ body: @differentiable (Element) -> Result
+  ) -> [Result] {
+    map(body)
+  }
+
+  @usableFromInline
+  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) = Swift.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/collection_hofs.swift

@@ -0,0 +1,36 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import StdlibUnittest
+#if os(macOS)
+import Darwin.C
+#else
+import Glibc
+#endif
+
+// Test suite for differentiable higher order functions for collections
+// such as `differentiableMap(_:)` and `differentiableReduce(_:)`.
+var CollectionHOFTests = TestSuite("CollectionHOF")
+
+let xx: [Float] = [1, 2, 3, 4, 5]
+
+CollectionHOFTests.test("differentiableMap(_:)") {
+  func double(_ xx: [Float]) -> [Float] {
+    xx.differentiableMap { $0 * $0 }
+  }
+  expectEqual([], pullback(at: xx, in: double)([]))
+  expectEqual([0], pullback(at: xx, in: double)([0]))
+  expectEqual([2], pullback(at: xx, in: double)([1]))
+  expectEqual([2, 4, 6, 8, 10], pullback(at: xx, in: double)([1, 1, 1, 1, 1]))
+}
+
+CollectionHOFTests.test("differentiableReduce(_:)") {
+  func product(_ xx: [Float]) -> Float {
+    xx.differentiableReduce(1) { $0 * $1 }
+  }
+  expectEqual([1], gradient(at: [0], in: product))
+  expectEqual([1], gradient(at: [1], in: product))
+  expectEqual([120, 60, 40, 30, 24], gradient(at: xx, in: product))
+}
+
+runAllTests()