From 39441de750ff26f0ee178433090963d88a428115 Mon Sep 17 00:00:00 2001 From: efremale Date: Thu, 9 Jul 2020 01:24:16 +0200 Subject: [PATCH 1/5] [AutoDiff] [stdlib] Add a JVP for 'differentiableMap(_:)'. --- .../ArrayDifferentiation.swift | 21 +++++++++++++++++++ .../validation-test/forward_mode.swift | 17 +++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/stdlib/public/Differentiation/ArrayDifferentiation.swift b/stdlib/public/Differentiation/ArrayDifferentiation.swift index a4b1149a20f56..a36b432bf25b0 100644 --- a/stdlib/public/Differentiation/ArrayDifferentiation.swift +++ b/stdlib/public/Differentiation/ArrayDifferentiation.swift @@ -312,6 +312,27 @@ extension Array where Element: Differentiable { } return (value: values, pullback: pullback) } + + @inlinable + @derivative(of: differentiableMap) + internal func _jvpDifferentiableMap( + _ body: @differentiable (Element) -> Result + ) -> ( + value: [Result], + differential: (Array.TangentVector) -> Array.TangentVector + ) { + var values: [Result] = [] + var differentials: [(Element.TangentVector) -> Result.TangentVector] = [] + for x in self { + let (y, df) = valueWithDifferential(at: x, in: body) + values.append(y) + differentials.append(df) + } + func differential(_ tans: Array.TangentVector) -> Array.TangentVector { + .init(zip(tans.base, differentials).map { tan, df in df(tan) }) + } + return (value: values, differential: differential) + } } extension Array where Element: Differentiable { diff --git a/test/AutoDiff/validation-test/forward_mode.swift b/test/AutoDiff/validation-test/forward_mode.swift index 22b049158c432..aa28c8243e642 100644 --- a/test/AutoDiff/validation-test/forward_mode.swift +++ b/test/AutoDiff/validation-test/forward_mode.swift @@ -1319,4 +1319,21 @@ ForwardModeTests.test("ForceUnwrapping") { expectEqual(5, forceUnwrap(Float(2))) } +//===----------------------------------------------------------------------===// +// Array methods from ArrayDifferentiation.swift +//===----------------------------------------------------------------------===// +ForwardModeTests.test("Array.differentiableMap") { + let a: Array = [1, 2, 3] + + func multiplyMap(a: Array) -> Array { + return a.differentiableMap({x in 3 * x}); + } + expectEqual([3, 3, 3], differential(at: [1, 1, 1], in: multiplyMap)(a)) + + func squareMap(a: Array) -> Array { + return a.differentiableMap({x in x * x}); + } + expectEqual([2, 4, 6], differential(at: [1, 1, 1], in: squareMap)(a)) +} + runAllTests() From 64fb5f95f8b6c39c3cca11c4547c70ef8f21d1df Mon Sep 17 00:00:00 2001 From: efremale Date: Fri, 10 Jul 2020 14:45:53 +0200 Subject: [PATCH 2/5] Fix code style --- .../public/Differentiation/ArrayDifferentiation.swift | 2 +- test/AutoDiff/validation-test/forward_mode.swift | 11 ++++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/stdlib/public/Differentiation/ArrayDifferentiation.swift b/stdlib/public/Differentiation/ArrayDifferentiation.swift index a36b432bf25b0..1f8cbf9b09f33 100644 --- a/stdlib/public/Differentiation/ArrayDifferentiation.swift +++ b/stdlib/public/Differentiation/ArrayDifferentiation.swift @@ -316,7 +316,7 @@ extension Array where Element: Differentiable { @inlinable @derivative(of: differentiableMap) internal func _jvpDifferentiableMap( - _ body: @differentiable (Element) -> Result + _ body: @differentiable (Element) -> Result ) -> ( value: [Result], differential: (Array.TangentVector) -> Array.TangentVector diff --git a/test/AutoDiff/validation-test/forward_mode.swift b/test/AutoDiff/validation-test/forward_mode.swift index aa28c8243e642..e3fa4ae572d0b 100644 --- a/test/AutoDiff/validation-test/forward_mode.swift +++ b/test/AutoDiff/validation-test/forward_mode.swift @@ -1322,16 +1322,17 @@ ForwardModeTests.test("ForceUnwrapping") { //===----------------------------------------------------------------------===// // Array methods from ArrayDifferentiation.swift //===----------------------------------------------------------------------===// + ForwardModeTests.test("Array.differentiableMap") { - let a: Array = [1, 2, 3] + let a: [Float] = [1, 2, 3] - func multiplyMap(a: Array) -> Array { - return a.differentiableMap({x in 3 * x}); + func multiplyMap(_ a: [Float]) -> [Float] { + return a.differentiableMap({ x in 3 * x }) } expectEqual([3, 3, 3], differential(at: [1, 1, 1], in: multiplyMap)(a)) - func squareMap(a: Array) -> Array { - return a.differentiableMap({x in x * x}); + func squareMap(_ a: [Float]) -> [Float] { + return a.differentiableMap({ x in x * x }) } expectEqual([2, 4, 6], differential(at: [1, 1, 1], in: squareMap)(a)) } From cb8ea0c8e81f7c7f3b250b8897ad8a0c6f66d06d Mon Sep 17 00:00:00 2001 From: efremale Date: Sat, 11 Jul 2020 00:53:30 +0200 Subject: [PATCH 3/5] Fix type error --- test/AutoDiff/validation-test/forward_mode.swift | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/AutoDiff/validation-test/forward_mode.swift b/test/AutoDiff/validation-test/forward_mode.swift index e3fa4ae572d0b..67cb49401d8f2 100644 --- a/test/AutoDiff/validation-test/forward_mode.swift +++ b/test/AutoDiff/validation-test/forward_mode.swift @@ -1325,16 +1325,17 @@ ForwardModeTests.test("ForceUnwrapping") { ForwardModeTests.test("Array.differentiableMap") { let a: [Float] = [1, 2, 3] + let tan = Array.TangentVector.init(a) func multiplyMap(_ a: [Float]) -> [Float] { return a.differentiableMap({ x in 3 * x }) } - expectEqual([3, 3, 3], differential(at: [1, 1, 1], in: multiplyMap)(a)) + expectEqual([3, 3, 3], differential(at: [1, 1, 1], in: multiplyMap)(tan)) func squareMap(_ a: [Float]) -> [Float] { return a.differentiableMap({ x in x * x }) } - expectEqual([2, 4, 6], differential(at: [1, 1, 1], in: squareMap)(a)) + expectEqual([2, 4, 6], differential(at: [1, 1, 1], in: squareMap)(tan)) } runAllTests() From 45201e5bf7bb22a7efd9688f07d4a9576b0de9d8 Mon Sep 17 00:00:00 2001 From: efremale Date: Sat, 11 Jul 2020 01:44:47 +0200 Subject: [PATCH 4/5] Fix differentiableMap test --- test/AutoDiff/validation-test/forward_mode.swift | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/AutoDiff/validation-test/forward_mode.swift b/test/AutoDiff/validation-test/forward_mode.swift index 67cb49401d8f2..d3c4280c6910e 100644 --- a/test/AutoDiff/validation-test/forward_mode.swift +++ b/test/AutoDiff/validation-test/forward_mode.swift @@ -1324,18 +1324,18 @@ ForwardModeTests.test("ForceUnwrapping") { //===----------------------------------------------------------------------===// ForwardModeTests.test("Array.differentiableMap") { - let a: [Float] = [1, 2, 3] - let tan = Array.TangentVector.init(a) + let x: [Float] = [1, 2, 3] + let tan = Array.TangentVector([1, 1, 1]) func multiplyMap(_ a: [Float]) -> [Float] { return a.differentiableMap({ x in 3 * x }) } - expectEqual([3, 3, 3], differential(at: [1, 1, 1], in: multiplyMap)(tan)) + expectEqual([3, 3, 3], differential(at: x, in: multiplyMap)(tan)) func squareMap(_ a: [Float]) -> [Float] { return a.differentiableMap({ x in x * x }) } - expectEqual([2, 4, 6], differential(at: [1, 1, 1], in: squareMap)(tan)) + expectEqual([2, 4, 6], differential(at: x, in: squareMap)(tan)) } runAllTests() From d31f3acaf46a10bc9df84b943e25e51a92c29513 Mon Sep 17 00:00:00 2001 From: efremale Date: Sat, 11 Jul 2020 20:41:27 +0200 Subject: [PATCH 5/5] [AutoDiff][stdlib] Add JVPs to ArrayDifferentiation.swift `_jvpDifferentiableReduce` is taken from https://github.com/apple/swift/pull/29324 --- .../ArrayDifferentiation.swift | 85 ++++++++++++++++++ .../validation-test/forward_mode.swift | 89 +++++++++++++++++++ 2 files changed, 174 insertions(+) diff --git a/stdlib/public/Differentiation/ArrayDifferentiation.swift b/stdlib/public/Differentiation/ArrayDifferentiation.swift index 1f8cbf9b09f33..177e9a3572d6e 100644 --- a/stdlib/public/Differentiation/ArrayDifferentiation.swift +++ b/stdlib/public/Differentiation/ArrayDifferentiation.swift @@ -42,6 +42,14 @@ where Element: Differentiable { return (base, { $0 }) } + @usableFromInline + @derivative(of: base) + func _jvpBase() -> ( + value: [Element], differential: (Array.TangentVector) -> TangentVector + ) { + return (base, { $0 }) + } + /// Creates a differentiable view of the given array. public init(_ base: [Element]) { self._base = base } @@ -53,6 +61,14 @@ where Element: Differentiable { return (Array.DifferentiableView(base), { $0 }) } + @usableFromInline + @derivative(of: init(_:)) + static func _jvpInit(_ base: [Element]) -> ( + value: Array.DifferentiableView, differential: (TangentVector) -> TangentVector + ) { + return (Array.DifferentiableView(base), { $0 }) + } + public typealias TangentVector = Array.DifferentiableView @@ -191,6 +207,17 @@ extension Array where Element: Differentiable { return (self[index], pullback) } + @usableFromInline + @derivative(of: subscript) + func _jvpSubscript(index: Int) -> ( + value: Element, differential: (TangentVector) -> Element.TangentVector + ) { + func differential(_ v: TangentVector) -> Element.TangentVector { + return v[index] + } + return (self[index], differential) + } + @usableFromInline @derivative(of: +) static func _vjpConcatenate(_ lhs: Self, _ rhs: Self) -> ( @@ -210,8 +237,26 @@ extension Array where Element: Differentiable { } return (lhs + rhs, pullback) } + + @usableFromInline + @derivative(of: +) + static func _jvpConcatenate(_ lhs: Self, _ rhs: Self) -> ( + value: Self, + differential: (TangentVector, TangentVector) -> TangentVector + ) { + func differential(_ l: TangentVector, _ r: TangentVector) -> TangentVector { + precondition( + l.base.count == lhs.count && r.base.count == rhs.count, """ + Tangent vectors with invalid count; expected to equal the \ + operand counts \(lhs.count) and \(rhs.count) + """) + return .init(l.base + r.base) + } + return (lhs + rhs, differential) + } } + extension Array where Element: Differentiable { @usableFromInline @derivative(of: append) @@ -277,6 +322,17 @@ extension Array where Element: Differentiable { } ) } + + @usableFromInline + @derivative(of: init(repeating:count:)) + static func _jvpInit(repeating repeatedValue: Element, count: Int) -> ( + value: Self, differential: (Element.TangentVector) -> TangentVector + ) { + ( + value: Self(repeating: repeatedValue, count: count), + differential: { v in TangentVector(.init(repeating: v, count: count)) } + ) + } } //===----------------------------------------------------------------------===// @@ -382,4 +438,33 @@ extension Array where Element: Differentiable { } ) } + + @inlinable + @derivative(of: differentiableReduce, wrt: (self, initialResult)) + func _jvpDifferentiableReduce( + _ initialResult: Result, + _ nextPartialResult: @differentiable (Result, Element) -> Result + ) -> (value: Result, + differential: (Array.TangentVector, Result.TangentVector) + -> Result.TangentVector) { + var differentials: + [(Result.TangentVector, Element.TangentVector) -> Result.TangentVector] + = [] + let count = self.count + differentials.reserveCapacity(count) + var result = initialResult + for element in self { + let (y, df) = + valueWithDifferential(at: result, element, in: nextPartialResult) + result = y + differentials.append(df) + } + return (value: result, differential: { dSelf, dInitial in + var dResult = dInitial + for (dElement, df) in zip(dSelf.base, differentials) { + dResult = df(dResult, dElement) + } + return dResult + }) + } } diff --git a/test/AutoDiff/validation-test/forward_mode.swift b/test/AutoDiff/validation-test/forward_mode.swift index d3c4280c6910e..e43740c1244a5 100644 --- a/test/AutoDiff/validation-test/forward_mode.swift +++ b/test/AutoDiff/validation-test/forward_mode.swift @@ -1323,6 +1323,75 @@ ForwardModeTests.test("ForceUnwrapping") { // Array methods from ArrayDifferentiation.swift //===----------------------------------------------------------------------===// +typealias FloatArrayTan = Array.TangentVector + +ForwardModeTests.test("Array.+") { + func sumFirstThreeConcatenating(_ a: [Float], _ b: [Float]) -> Float { + let c = a + b + return c[0] + c[1] + c[2] + } + + expectEqual(3, differential(at: [0, 0], [0, 0], in: sumFirstThreeConcatenating)(.init([1, 1]), .init([1, 1]))) + expectEqual(0, differential(at: [0, 0], [0, 0], in: sumFirstThreeConcatenating)(.init([0, 0]), .init([0, 1]))) + expectEqual(1, differential(at: [0, 0], [0, 0], in: sumFirstThreeConcatenating)(.init([0, 1]), .init([0, 1]))) + expectEqual(1, differential(at: [0, 0], [0, 0], in: sumFirstThreeConcatenating)(.init([1, 0]), .init([0, 1]))) + expectEqual(1, differential(at: [0, 0], [0, 0], in: sumFirstThreeConcatenating)(.init([0, 0]), .init([1, 1]))) + expectEqual(2, differential(at: [0, 0], [0, 0], in: sumFirstThreeConcatenating)(.init([1, 1]), .init([0, 1]))) + + expectEqual( + 3, + differential(at: [0, 0, 0, 0], [0, 0], in: sumFirstThreeConcatenating)(.init([1, 1, 1, 1]), .init([1, 1]))) + expectEqual( + 3, + differential(at: [0, 0, 0, 0], [0, 0], in: sumFirstThreeConcatenating)(.init([1, 1, 1, 0]), .init([0, 0]))) + + expectEqual( + 3, + differential(at: [], [0, 0, 0, 0], in: sumFirstThreeConcatenating)(.init([]), .init([1, 1, 1, 1]))) + expectEqual( + 0, + differential(at: [], [0, 0, 0, 0], in: sumFirstThreeConcatenating)(.init([]), .init([0, 0, 0, 1]))) +} + +ForwardModeTests.test("Array.init(repeating:count:)") { + @differentiable + func repeating(_ x: Float) -> [Float] { + Array(repeating: x, count: 10) + } + expectEqual(Float(10), derivative(at: .zero) { x in + repeating(x).differentiableReduce(0, {$0 + $1}) + }) + expectEqual(Float(20), differential(at: .zero, in: { x in + repeating(x).differentiableReduce(0, {$0 + $1}) + })(2)) +} + +ForwardModeTests.test("Array.DifferentiableView.init") { + @differentiable + func constructView(_ x: [Float]) -> Array.DifferentiableView { + return Array.DifferentiableView(x) + } + + let forward = differential(at: [5, 6, 7, 8], in: constructView) + expectEqual( + FloatArrayTan([1, 2, 3, 4]), + forward(FloatArrayTan([1, 2, 3, 4]))) +} + +ForwardModeTests.test("Array.DifferentiableView.base") { + @differentiable + func accessBase(_ x: Array.DifferentiableView) -> [Float] { + return x.base + } + + let forward = differential( + at: Array.DifferentiableView([5, 6, 7, 8]), + in: accessBase) + expectEqual( + FloatArrayTan([1, 2, 3, 4]), + forward(FloatArrayTan([1, 2, 3, 4]))) +} + ForwardModeTests.test("Array.differentiableMap") { let x: [Float] = [1, 2, 3] let tan = Array.TangentVector([1, 1, 1]) @@ -1338,4 +1407,24 @@ ForwardModeTests.test("Array.differentiableMap") { expectEqual([2, 4, 6], differential(at: x, in: squareMap)(tan)) } +ForwardModeTests.test("Array.differentiableReduce") { + let x: [Float] = [1, 2, 3] + let tan = Array.TangentVector([1, 1, 1]) + + func sumReduce(_ a: [Float]) -> Float { + return a.differentiableReduce(0, { $0 + $1 }) + } + expectEqual(1 + 1 + 1, differential(at: x, in: sumReduce)(tan)) + + func productReduce(_ a: [Float]) -> Float { + return a.differentiableReduce(1, { $0 * $1 }) + } + expectEqual(x[1] * x[2] + x[0] * x[2] + x[0] * x[1], differential(at: x, in: productReduce)(tan)) + + func sumOfSquaresReduce(_ a: [Float]) -> Float { + return a.differentiableReduce(0, { $0 + $1 * $1 }) + } + expectEqual(2 * x[0] + 2 * x[1] + 2 * x[2], differential(at: x, in: sumOfSquaresReduce)(tan)) +} + runAllTests()