-
Notifications
You must be signed in to change notification settings - Fork 10.3k
/
differentiable_attr.swift
162 lines (145 loc) · 5.8 KB
/
differentiable_attr.swift
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
// SWIFT_ENABLE_TENSORFLOW
// RUN: %empty-directory(%t)
// RUN: %target-swift-frontend %s -emit-module -parse-as-library -o %t
// RUN: %target-sil-opt -disable-sil-linking -enable-sil-verify-all %t/differentiable_attr.swiftmodule -o - | %FileCheck %s
struct CheckpointsFoo {}
func pfoo(_ x: Float) -> (checkpoints: CheckpointsFoo, originalValue: Float) {
return (CheckpointsFoo(), x * x)
}
func dfoo_checkpointed(_ seed: Float, _ checkpoints: CheckpointsFoo, _ originalValue: Float, _ x: Float) -> Float {
return 2 * x
}
// CHECK-DAG: @differentiable(primal: pfoo, adjoint: dfoo_checkpointed)
// CHECK-DAG: func foo_checkpointed(_ x: Float) -> Float
@differentiable(primal: pfoo(_:), adjoint: dfoo_checkpointed)
func foo_checkpointed(_ x: Float) -> Float {
return x * x
}
struct S<T> {
struct Checkpoints {
let s: S
}
func primal(x: Float) -> (Checkpoints, Float) {
return (Checkpoints(s: self), x)
}
func adjoint_checkpointed(_: Float, _: Checkpoints, _: Float, _ x: Float) -> S {
return self
}
// CHECK-DAG: @differentiable(wrt: (self), primal: primal, adjoint: adjoint_checkpointed)
// CHECK-DAG: func original(x: Float) -> Float
@differentiable(wrt: (self), primal: primal, adjoint: adjoint_checkpointed)
func original(x: Float) -> Float {
return x
}
}
extension S : Differentiable, VectorNumeric {
typealias TangentVector = S
typealias CotangentVector = S
typealias Scalar = Float
static var zero: S { fatalError("unimplemented") }
static func + (lhs: S, rhs: S) -> S { fatalError("unimplemented") }
static func - (lhs: S, rhs: S) -> S { fatalError("unimplemented") }
static func * (lhs: Float, rhs: S) -> S { fatalError("unimplemented") }
}
func pbaz1<T : Differentiable>(_ x: T, _ y: T) -> ((T, T), T) {
return ((y, y), x)
}
func dbaz1_checkpointed<T : Differentiable>(_ seed: T.CotangentVector, _ primal: (T, T), _ originalValue: T, _ x: T, _ y: T) -> (T.CotangentVector, T.CotangentVector) {
return (seed, seed)
}
// CHECK-DAG: @differentiable(primal: pbaz1, adjoint: dbaz1_checkpointed)
// CHECK-DAG: func baz1_checkpointed<T>(_ x: T, _ y: T) -> T
@differentiable(primal: pbaz1(_:_:), adjoint: dbaz1_checkpointed)
func baz1_checkpointed<T : Differentiable>(_ x: T, _ y: T) -> T {
return x
}
struct CheckpointsFP<T : FloatingPoint> {
let meow: T
}
func pbaz2<T : Differentiable & FloatingPoint>(_ x: T, _ y: T) -> (CheckpointsFP<T>, T) {
return (CheckpointsFP(meow: 1), x + y)
}
func dbaz2_checkpointed<T : Differentiable & FloatingPoint>(_ seed: T.CotangentVector, _ primal: CheckpointsFP<T>, _ originalValue: T, _ x: T, _ y: T) -> (T.CotangentVector, T.CotangentVector) {
return (seed, seed)
}
// CHECK-DAG: @differentiable(primal: pbaz2, adjoint: dbaz2_checkpointed)
// CHECK-DAG: func baz2_checkpointed<T>(_ x: T, _ y: T) -> T where T : Differentiable, T : FloatingPoint
@differentiable(primal: pbaz2(_:_:), adjoint: dbaz2_checkpointed)
func baz2_checkpointed<T : Differentiable & FloatingPoint>(_ x: T, _ y: T) -> T {
return x
}
@differentiable(jvp: jvpSimpleJVP)
func jvpSimple(x: Float) -> Float {
return x
}
// CHECK-DAG: @differentiable(jvp: jvpSimpleJVP)
// CHECK-DAG: func jvpSimpleJVP(x: Float) -> (Float, (Float) -> Float)
func jvpSimpleJVP(x: Float) -> (Float, (Float) -> Float) {
return (x, { v in v })
}
@differentiable(vjp: vjpSimpleVJP)
func vjpSimple(x: Float) -> Float {
return x
}
// CHECK-DAG: @differentiable(vjp: vjpSimpleVJP)
// CHECK-DAG: func vjpSimpleVJP(x: Float) -> (Float, (Float) -> Float)
func vjpSimpleVJP(x: Float) -> (Float, (Float) -> Float) {
return (x, { v in v })
}
/*
// FIXME: 'where' clauses don't show up after round-trip serialization/deserialization.
// HECK-DAG: @differentiable(vjp: vjpTestWhereClause where T : Differentiable)
// HECK-DAG: func testWhereClause<T : Numeric>(x: T) -> T {
@differentiable(vjp: vjpTestWhereClause where T : Differentiable)
func testWhereClause<T : Numeric>(x: T) -> T {
return x
}
func vjpTestWhereClause<T>(x: T) -> (T, (T.CotangentVector) -> T.CotangentVector)
where T : Numeric, T : Differentiable
{
return (x, { v in v })
}
protocol P {}
extension P {
// HECK-DAG: @differentiable(wrt: (self), vjp: vjpTestWhereClause where Self : Differentiable)
// HECK-DAG: func testWhereClause() -> Self {
@differentiable(wrt: (self), vjp: vjpTestWhereClause where Self : Differentiable)
func testWhereClause() -> Self {
return self
}
}
extension P where Self : Differentiable {
func vjpTestWhereClause() -> (Self, (Self.CotangentVector) -> Self.CotangentVector) {
return (self, { v in v })
}
}
*/
/*
// NOTE: The failing tests involve where clauses with member type constraints.
// They pass type-checking but crash during serialization.
// HECK-DAG: @differentiable(vjp: vjpTestWhereClauseMemberTypeConstraint where T : Differentiable, T == T.CotangentVector)
// HECK-DAG: func testWhereClauseMemberTypeConstraint<T : Numeric>(x: T) -> T {
@differentiable(vjp: vjpTestWhereClauseMemberTypeConstraint where T : Differentiable, T == T.CotangentVector)
func testWhereClauseMemberTypeConstraint<T : Numeric>(x: T) -> T {
return x
}
func vjpTestWhereClauseMemberTypeConstraint<T>(x: T) -> (T, (T) -> T)
where T : Numeric, T : Differentiable, T == T.CotangentVector
{
return (x, { v in v })
}
protocol P {}
extension P {
// HECK-DAG: @differentiable(wrt: (self), vjp: vjpTestWhereClauseMemberTypeConstraint where Self.CotangentVector == Self, Self : Differentiable)
// HECK-DAG: func testWhereClauseMemberTypeConstraint() -> Self {
@differentiable(wrt: (self), vjp: vjpTestWhereClauseMemberTypeConstraint where Self.CotangentVector == Self, Self : Differentiable)
func testWhereClauseMemberTypeConstraint() -> Self {
return self
}
}
extension P where Self : Differentiable {
func vjpTestWhereClauseMemberTypeConstraint() -> (Self, (Self.CotangentVector) -> Self.CotangentVector) {
return (self, { v in v })
}
}
*/