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 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251
|
// RUN: %target-run-simple-swift
// REQUIRES: executable_test
// Would fail due to unavailability of swift_autoDiffCreateLinearMapContext.
// `inout` parameter differentiation tests.
import DifferentiationUnittest
import StdlibUnittest
var InoutParameterAutoDiffTests = TestSuite("InoutParameterDifferentiation")
// TODO(TF-1173): Move floating-point mutating operation tests to
// `test/AutoDiff/stdlib/floating_point.swift.gyb` when forward-mode
// differentiation supports `inout` parameter differentiation.
InoutParameterAutoDiffTests.test("Float.+=") {
func mutatingAddWrapper(_ x: Float, _ y: Float) -> Float {
var result: Float = x
result += y
return result
}
expectEqual((1, 1), gradient(at: 4, 5, of: mutatingAddWrapper))
expectEqual((10, 10), pullback(at: 4, 5, of: mutatingAddWrapper)(10))
}
InoutParameterAutoDiffTests.test("Float.-=") {
func mutatingSubtractWrapper(_ x: Float, _ y: Float) -> Float {
var result: Float = x
result += y
return result
}
expectEqual((1, 1), gradient(at: 4, 5, of: mutatingSubtractWrapper))
expectEqual((10, 10), pullback(at: 4, 5, of: mutatingSubtractWrapper)(10))
}
InoutParameterAutoDiffTests.test("Float.*=") {
func mutatingMultiplyWrapper(_ x: Float, _ y: Float) -> Float {
var result: Float = x
result += y
return result
}
expectEqual((1, 1), gradient(at: 4, 5, of: mutatingMultiplyWrapper))
expectEqual((10, 10), pullback(at: 4, 5, of: mutatingMultiplyWrapper)(10))
}
InoutParameterAutoDiffTests.test("Float./=") {
func mutatingDivideWrapper(_ x: Float, _ y: Float) -> Float {
var result: Float = x
result += y
return result
}
expectEqual((1, 1), gradient(at: 4, 5, of: mutatingDivideWrapper))
expectEqual((10, 10), pullback(at: 4, 5, of: mutatingDivideWrapper)(10))
}
// Simplest possible `inout` parameter differentiation.
InoutParameterAutoDiffTests.test("InoutIdentity") {
// Semantically, an empty function with an `inout` parameter is an identity
// function.
func inoutIdentity(_ x: inout Float) {}
func identity(_ x: Float) -> Float {
var result = x
inoutIdentity(&result)
return result
}
expectEqual(1, gradient(at: 10, of: identity))
expectEqual(10, pullback(at: 10, of: identity)(10))
}
extension Float {
// Custom version of `Float.*=`, implemented using `Float.*` and mutation.
// Verify that its generated derivative has the same behavior as the
// registered derivative for `Float.*=`.
@differentiable(reverse)
static func multiplyAssign(_ lhs: inout Float, _ rhs: Float) {
lhs = lhs * rhs
}
}
InoutParameterAutoDiffTests.test("ControlFlow") {
func sum(_ array: [Float]) -> Float {
var result: Float = 0
for i in withoutDerivative(at: array.indices) {
result += array[i]
}
return result
}
expectEqual([1, 1, 1], gradient(at: [1, 2, 3], of: sum))
func product(_ array: [Float]) -> Float {
var result: Float = 1
for i in withoutDerivative(at: array.indices) {
result *= array[i]
}
return result
}
expectEqual([20, 15, 12], gradient(at: [3, 4, 5], of: product))
func productCustom(_ array: [Float]) -> Float {
var result: Float = 1
for i in withoutDerivative(at: array.indices) {
Float.multiplyAssign(&result, array[i])
}
return result
}
expectEqual([20, 15, 12], gradient(at: [3, 4, 5], of: productCustom))
}
InoutParameterAutoDiffTests.test("SetAccessor") {
struct S: Differentiable {
var x: Float
var computed: Float {
get { x }
set { x = newValue }
}
// Computed property with explicit `@differentiable` accessors.
var doubled: Float {
@differentiable(reverse)
get { x + x }
@differentiable(reverse)
set { x = newValue / 2 }
}
}
// `squared` implemented using a `set` accessor.
func squared(_ x: Float) -> Float {
var s = S(x: 1)
s.x *= x
s.computed *= x
return s.x
}
expectEqual((9, 6), valueWithGradient(at: 3, of: squared))
expectEqual((16, 8), valueWithGradient(at: 4, of: squared))
// `quadrupled` implemented using a `set` accessor.
func quadrupled(_ x: Float) -> Float {
var s = S(x: 1)
s.doubled *= 4 * x
return s.x
}
print(valueWithGradient(at: 3, of: quadrupled))
print(valueWithGradient(at: 4, of: quadrupled))
expectEqual((12, 4), valueWithGradient(at: 3, of: quadrupled))
expectEqual((16, 4), valueWithGradient(at: 4, of: quadrupled))
}
// Test differentiation wrt `inout` parameters that have a class type.
InoutParameterAutoDiffTests.test("InoutClassParameter") {
class Class: Differentiable {
@differentiable(reverse)
var x: Float
init(_ x: Float) {
self.x = x
}
}
do {
func squaredViaMutation(_ c: inout Class) {
c = Class(c.x * c.x)
}
func squared(_ x: Float) -> Float {
var c = Class(x)
squaredViaMutation(&c)
return c.x
}
expectEqual((100, 20), valueWithGradient(at: 10, of: squared))
expectEqual(200, pullback(at: 10, of: squared)(10))
}
do {
func squaredViaModifyAccessor(_ c: inout Class) {
// The line below calls `Class.x.modify`.
c.x *= c.x
}
func squared(_ x: Float) -> Float {
var c = Class(x)
squaredViaModifyAccessor(&c)
return c.x
}
// FIXME(TF-1080): Fix incorrect class property `modify` accessor derivative values.
// expectEqual((100, 20), valueWithGradient(at: 10, of: squared))
// expectEqual(200, pullback(at: 10, of: squared)(10))
expectEqual((100, 1), valueWithGradient(at: 10, of: squared))
expectEqual(10, pullback(at: 10, of: squared)(10))
}
}
// Test function with wrt `inout` parameter, which should be treated as a differentiability result.
// Original issue https://github.com/apple/swift/issues/55745 deals with non-wrt `inout` which
// we explicitly disallow now
protocol P_55745 {
@differentiable(reverse, wrt: (x, y))
func method(_ x: Float, _ y: inout Float)
@differentiable(reverse, wrt: (x, y))
func genericMethod<T: Differentiable>(_ x: T, _ y: inout T)
}
InoutParameterAutoDiffTests.test("non-wrt inout parameter") {
struct Struct: P_55745 {
@differentiable(reverse, wrt: (x, y))
func method(_ x: Float, _ y: inout Float) {
y = y * x
}
@differentiable(reverse, wrt: (x, y))
func genericMethod<T: Differentiable>(_ x: T, _ y: inout T) {
y = x
}
}
@differentiable(reverse, wrt: x)
func foo(_ s: Struct, _ x: Float, _ y: Float) -> Float {
var y = y
s.method(x, &y)
return y
}
@differentiable(reverse, wrt: x)
func fooGeneric<T: P_55745>(_ s: T, _ x: Float, _ y: Float) -> Float {
var y = y
s.method(x, &y)
return x
}
let s = Struct()
do {
let (value, (dx, dy)) = valueWithGradient(at: 2, 3, of: { foo(s, $0, $1) })
expectEqual(6, value)
expectEqual((3, 2), (dx, dy))
}
expectEqual((value: 6, gradient: 3), valueWithGradient(at: 2, of: { foo(s, $0, 3) }))
do {
let (value, (dx, dy)) = valueWithGradient(at: 2, 3, of: { fooGeneric(s, $0, $1) })
expectEqual(2, value)
expectEqual((1, 0), (dx, dy))
}
expectEqual((value: 2, gradient: 1), valueWithGradient(at: 2, of: { fooGeneric(s, $0, 3) }))
}
runAllTests()
|