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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
|
// RUN: %target-run-simple-swift
// REQUIRES: executable_test
// An end-to-end test that we can differentiate property accesses, with custom
// VJPs for the properties specified in various ways.
import StdlibUnittest
import DifferentiationUnittest
var E2EDifferentiablePropertyTests = TestSuite("E2EDifferentiableProperty")
struct TangentSpace : AdditiveArithmetic {
let x, y: Tracked<Float>
}
extension TangentSpace : Differentiable {
typealias TangentVector = TangentSpace
}
struct Space {
/// `x` is a computed property with a custom vjp.
var x: Tracked<Float> {
@differentiable(reverse)
get { storedX }
set { storedX = newValue }
}
@derivative(of: x)
func vjpX() -> (value: Tracked<Float>, pullback: (Tracked<Float>) -> TangentSpace) {
return (x, { v in TangentSpace(x: v, y: 0) } )
}
private var storedX: Tracked<Float>
@differentiable(reverse)
var y: Tracked<Float>
init(x: Tracked<Float>, y: Tracked<Float>) {
self.storedX = x
self.y = y
}
}
extension Space : Differentiable {
typealias TangentVector = TangentSpace
mutating func move(by offset: TangentSpace) {
x.move(by: offset.x)
y.move(by: offset.y)
}
}
E2EDifferentiablePropertyTests.testWithLeakChecking("computed property") {
let actualGrad = gradient(at: Space(x: 0, y: 0)) { (point: Space) -> Tracked<Float> in
return 2 * point.x
}
let expectedGrad = TangentSpace(x: 2, y: 0)
expectEqual(expectedGrad, actualGrad)
}
E2EDifferentiablePropertyTests.testWithLeakChecking("stored property") {
let actualGrad = gradient(at: Space(x: 0, y: 0)) { (point: Space) -> Tracked<Float> in
return 3 * point.y
}
let expectedGrad = TangentSpace(x: 0, y: 3)
expectEqual(expectedGrad, actualGrad)
}
struct GenericMemberWrapper<T : Differentiable> : Differentiable {
// Stored property.
@differentiable(reverse)
var x: T
func vjpX() -> (T, (T.TangentVector) -> GenericMemberWrapper.TangentVector) {
return (x, { TangentVector(x: $0) })
}
}
E2EDifferentiablePropertyTests.testWithLeakChecking("generic stored property") {
let actualGrad = gradient(at: GenericMemberWrapper<Tracked<Float>>(x: 1)) { point in
return 2 * point.x
}
let expectedGrad = GenericMemberWrapper<Tracked<Float>>.TangentVector(x: 2)
expectEqual(expectedGrad, actualGrad)
}
struct ProductSpaceSelfTangent : AdditiveArithmetic {
let x, y: Tracked<Float>
}
extension ProductSpaceSelfTangent : Differentiable {
typealias TangentVector = ProductSpaceSelfTangent
}
E2EDifferentiablePropertyTests.testWithLeakChecking("fieldwise product space, self tangent") {
let actualGrad = gradient(at: ProductSpaceSelfTangent(x: 0, y: 0)) { (point: ProductSpaceSelfTangent) -> Tracked<Float> in
return 5 * point.y
}
let expectedGrad = ProductSpaceSelfTangent(x: 0, y: 5)
expectEqual(expectedGrad, actualGrad)
}
struct ProductSpaceOtherTangentTangentSpace : AdditiveArithmetic {
let x, y: Tracked<Float>
}
extension ProductSpaceOtherTangentTangentSpace : Differentiable {
typealias TangentVector = ProductSpaceOtherTangentTangentSpace
}
struct ProductSpaceOtherTangent {
var x, y: Tracked<Float>
}
extension ProductSpaceOtherTangent : Differentiable {
typealias TangentVector = ProductSpaceOtherTangentTangentSpace
mutating func move(by offset: ProductSpaceOtherTangentTangentSpace) {
x.move(by: offset.x)
y.move(by: offset.y)
}
}
E2EDifferentiablePropertyTests.testWithLeakChecking("fieldwise product space, other tangent") {
let actualGrad = gradient(
at: ProductSpaceOtherTangent(x: 0, y: 0)
) { (point: ProductSpaceOtherTangent) -> Tracked<Float> in
return 7 * point.y
}
let expectedGrad = ProductSpaceOtherTangentTangentSpace(x: 0, y: 7)
expectEqual(expectedGrad, actualGrad)
}
E2EDifferentiablePropertyTests.testWithLeakChecking("computed property") {
struct TF_544 : Differentiable {
var value: Tracked<Float>
@differentiable(reverse)
var computed: Tracked<Float> {
get { value }
set { value = newValue }
}
}
let actualGrad = gradient(at: TF_544(value: 2.4)) { x in
return x.computed * x.computed
}
let expectedGrad = TF_544.TangentVector(value: 4.8)
expectEqual(expectedGrad, actualGrad)
}
runAllTests()
|
