1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
|
// RUN: %target-swift-frontend -emit-sil -O -g %s | %FileCheck %s
// REQUIRES: swift_in_compiler
// Fix for https://github.com/apple/swift/issues/62608
// We need to emit separate debug info location for different adjoint buffers
// created for the single input variable
import _Differentiation
public extension Array {
@inlinable
@differentiable(reverse)
mutating func update(at index: Int, byCalling closure: @differentiable(reverse) (inout Element) -> Void) where Element: Differentiable {
closure(&self[index])
}
}
public func valueWithPullback<T>(
at x: T, of f: @differentiable(reverse) (inout T) -> Void
) -> (value: Void, pullback: (inout T.TangentVector) -> Void) {
@differentiable(reverse)
func nonInoutWrappingFunction(_ t: T) -> T {
var t = t
f(&t)
return t
}
let nonInoutPullback = pullback(at: x, of: nonInoutWrappingFunction)
return ((), { $0 = nonInoutPullback($0) })
}
@inlinable
public func pullback<T>(
at x: T, of f: @differentiable(reverse) (inout T) -> Void
) -> (inout T.TangentVector) -> Void {
return valueWithPullback(at: x, of: f).pullback
}
// CHECK-LABEL: sil private @$s4main19testUpdateByCallingyyKF8fOfArrayL_5arraySdSaySdG_tFySdzcfU_TJpSUpSr :
// CHECK: alloc_stack $Double, var, name "derivative of 'element' in scope at {{.*}} (scope #3)"
// CHECK: debug_value %{{.*}} : $Builtin.FPIEEE64, var, (name "derivative of 'element' in scope at {{.*}} (scope #1)"
public extension Array where Element: Differentiable {
@inlinable
@derivative(of: update(at:byCalling:))
mutating func vjpUpdate(
at index: Int,
byCalling closure: @differentiable(reverse) (inout Element) -> Void
)
->
(value: Void, pullback: (inout Self.TangentVector) -> Void)
{
let closurePullback = pullback(at: self[index], of: closure)
return (value: (), pullback: { closurePullback(&$0.base[index]) })
}
}
func testUpdateByCalling() throws {
@differentiable(reverse)
func fOfArray(array: [Double]) -> Double {
var array = array
var result = 0.0
for i in withoutDerivative(at: 0 ..< array.count) {
array.update(at: i, byCalling: { (element: inout Double) in
let initialElement = element
for _ in withoutDerivative(at: 0 ..< i) {
element *= initialElement
}
})
result += array[i]
}
return result
}
let array = [Double](repeating: 1.0, count: 3)
let expectedGradientOfFOfArray = [1.0, 2.0, 3.0]
let obtainedGradientOfFOfArray = gradient(at: array, of: fOfArray).base
}
|
