From 39441de750ff26f0ee178433090963d88a428115 Mon Sep 17 00:00:00 2001 From: efremale Date: Thu, 9 Jul 2020 01:24:16 +0200 Subject: [PATCH 1/4] [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/4] 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/4] 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/4] 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()