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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
|
// RUN: %target-run-simple-swift
// REQUIRES: executable_test
import _Differentiation
import StdlibUnittest
var ReabstractionE2ETests = TestSuite("ReabstractionE2E")
ReabstractionE2ETests.test("diff param concrete => generic") {
func grad<T: Differentiable>(_ f: @differentiable(reverse) (T) -> Float, at x: T) -> T.TangentVector {
gradient(at: x, of: f)
}
func inner(_ x: Float) -> Float {
7 * x * x
}
expectEqual(Float(7 * 2 * 3), grad(inner, at: 3))
}
ReabstractionE2ETests.test("nondiff param concrete => generic") {
func grad<T: Differentiable>(_ f: @differentiable(reverse) (Float, @noDerivative T) -> Float, at x: Float, _ y: T) -> Float {
gradient(at: x) { f($0, y) }
}
func inner(_ x: Float, _ y: Float) -> Float {
7 * x * x + y
}
expectEqual(Float(7 * 2 * 3), grad(inner, at: 3, 10))
}
ReabstractionE2ETests.test("diff param and nondiff param concrete => generic") {
func grad<T: Differentiable>(_ f: @differentiable(reverse) (T, @noDerivative T) -> Float, at x: T, _ y: T) -> T.TangentVector {
gradient(at: x) { f($0, y) }
}
func inner(_ x: Float, _ y: Float) -> Float {
7 * x * x + y
}
expectEqual(Float(7 * 2 * 3), grad(inner, at: 3, 10))
}
ReabstractionE2ETests.test("result concrete => generic") {
func grad<T: Differentiable>(_ f: @differentiable(reverse) (Float) -> T, at x: Float, seed: T.TangentVector) -> Float {
pullback(at: x, of: f)(seed)
}
func inner(_ x: Float) -> Float {
7 * x * x
}
expectEqual(Float(7 * 2 * 3), grad(inner, at: 3, seed: 1))
}
protocol HasFloat: Differentiable {
@differentiable(reverse)
var float: Float { get }
@differentiable(reverse)
init(float: Float)
}
extension Float: HasFloat {
@differentiable(reverse)
var float: Float { self }
@differentiable(reverse)
init(float: Float) { self = float }
}
ReabstractionE2ETests.test("diff param generic => concrete") {
func inner<T: HasFloat>(x: T) -> Float {
7 * x.float * x.float
}
let transformed: @differentiable(reverse) (Float) -> Float = inner
expectEqual(Float(7 * 3 * 3), transformed(3))
expectEqual(Float(7 * 2 * 3), gradient(at: 3, of: transformed))
}
ReabstractionE2ETests.test("nondiff param generic => concrete") {
func inner<T: HasFloat>(x: Float, y: T) -> Float {
7 * x * x + y.float
}
let transformed: @differentiable(reverse) (Float, @noDerivative Float) -> Float = inner
expectEqual(Float(7 * 3 * 3 + 10), transformed(3, 10))
expectEqual(Float(7 * 2 * 3), gradient(at: 3) { transformed($0, 10) })
}
ReabstractionE2ETests.test("diff param and nondiff param generic => concrete") {
func inner<T: HasFloat>(x: T, y: T) -> Float {
7 * x.float * x.float + y.float
}
let transformed: @differentiable(reverse) (Float, @noDerivative Float) -> Float = inner
expectEqual(Float(7 * 3 * 3 + 10), transformed(3, 10))
expectEqual(Float(7 * 2 * 3), gradient(at: 3) { transformed($0, 10) })
}
ReabstractionE2ETests.test("result generic => concrete") {
func inner<T: HasFloat>(x: Float) -> T {
T(float: 7 * x * x)
}
let transformed: @differentiable(reverse) (Float) -> Float = inner
expectEqual(Float(7 * 3 * 3), transformed(3))
expectEqual(Float(7 * 2 * 3), gradient(at: 3, of: transformed))
}
ReabstractionE2ETests.test("diff param concrete => generic => concrete") {
typealias FnTy<T: Differentiable> = @differentiable(reverse) (T) -> Float
func id<T: Differentiable>(_ f: @escaping FnTy<T>) -> FnTy<T> { f }
func inner(_ x: Float) -> Float {
7 * x * x
}
let transformed = id(inner)
expectEqual(Float(7 * 3 * 3), transformed(3))
expectEqual(Float(7 * 2 * 3), gradient(at: 3, of: transformed))
}
ReabstractionE2ETests.test("nondiff param concrete => generic => concrete") {
typealias FnTy<T: Differentiable> = @differentiable(reverse) (Float, @noDerivative T) -> Float
func id<T: Differentiable>(_ f: @escaping FnTy<T>) -> FnTy<T> { f }
func inner(_ x: Float, _ y: Float) -> Float {
7 * x * x + y
}
let transformed = id(inner)
expectEqual(Float(7 * 3 * 3 + 10), transformed(3, 10))
expectEqual(Float(7 * 2 * 3), gradient(at: 3) { transformed($0, 10) })
}
ReabstractionE2ETests.test("diff param and nondiff param concrete => generic => concrete") {
typealias FnTy<T: Differentiable> = @differentiable(reverse) (T, @noDerivative T) -> Float
func id<T: Differentiable>(_ f: @escaping FnTy<T>) -> FnTy<T> { f }
func inner(_ x: Float, _ y: Float) -> Float {
7 * x * x + y
}
let transformed = id(inner)
expectEqual(Float(7 * 3 * 3 + 10), transformed(3, 10))
expectEqual(Float(7 * 2 * 3), gradient(at: 3) { transformed($0, 10) })
}
ReabstractionE2ETests.test("result concrete => generic => concrete") {
typealias FnTy<T: Differentiable> = @differentiable(reverse) (Float) -> T
func id<T: Differentiable>(_ f: @escaping FnTy<T>) -> FnTy<T> { f }
func inner(_ x: Float) -> Float {
7 * x * x
}
let transformed = id(inner)
expectEqual(Float(7 * 3 * 3), transformed(3))
expectEqual(Float(7 * 2 * 3), gradient(at: 3, of: transformed))
}
ReabstractionE2ETests.test("@differentiable(reverse) function => opaque generic => concrete") {
func id<T>(_ t: T) -> T { t }
let inner: @differentiable(reverse) (Float) -> Float = { 7 * $0 * $0 }
let transformed = id(inner)
expectEqual(Float(7 * 3 * 3), transformed(3))
expectEqual(Float(7 * 2 * 3), gradient(at: 3, of: id(inner)))
}
ReabstractionE2ETests.test("@differentiable(reverse) function => opaque Any => concrete") {
func id(_ any: Any) -> Any { any }
let inner: @differentiable(reverse) (Float) -> Float = { 7 * $0 * $0 }
let transformed = id(inner)
let casted = transformed as! @differentiable(reverse) (Float) -> Float
expectEqual(Float(7 * 3 * 3), casted(3))
expectEqual(Float(7 * 2 * 3), gradient(at: 3, of: casted))
}
ReabstractionE2ETests.test("access @differentiable(reverse) function using KeyPath") {
struct Container {
let f: @differentiable(reverse) (Float) -> Float
}
let container = Container(f: { 7 * $0 * $0 })
let kp = \Container.f
let extracted = container[keyPath: kp]
expectEqual(Float(7 * 3 * 3), extracted(3))
expectEqual(Float(7 * 2 * 3), gradient(at: 3, of: extracted))
}
runAllTests()
|
