| // RUN: %target-run-simple-swift |
| |
| // NOTE(TF-813): verify that enabling forward-mode does not affect reverse-mode. |
| // Temporarily disabled because forward-mode is not at feature parity with reverse-mode. |
| // UN: %target-run-simple-swift(-Xfrontend -enable-experimental-forward-mode-differentiation) |
| |
| // RUN: %target-swift-frontend -Xllvm -sil-print-after=differentiation %s -emit-sil -o /dev/null -module-name null 2>&1 | %FileCheck %s |
| |
| // REQUIRES: executable_test |
| |
| import StdlibUnittest |
| import DifferentiationUnittest |
| |
| var SimpleMathTests = TestSuite("SimpleMath") |
| |
| SimpleMathTests.test("Arithmetics") { |
| func foo1(x: Float, y: Float) -> Float { |
| return x * y |
| } |
| expectEqual((4, 3), gradient(at: 3, 4, in: foo1)) |
| func foo2(x: Float, y: Float) -> Float { |
| return -x * y |
| } |
| expectEqual((-4, -3), gradient(at: 3, 4, in: foo2)) |
| func foo3(x: Float, y: Float) -> Float { |
| return -x + y |
| } |
| expectEqual((-1, 1), gradient(at: 3, 4, in: foo3)) |
| } |
| |
| SimpleMathTests.test("Fanout") { |
| func foo1(x: Float) -> Float { |
| x - x |
| } |
| expectEqual(0, gradient(at: 100, in: foo1)) |
| func foo2(x: Float) -> Float { |
| x + x |
| } |
| expectEqual(2, gradient(at: 100, in: foo2)) |
| func foo3(x: Float, y: Float) -> Float { |
| x + x + x * y |
| } |
| expectEqual((4, 3), gradient(at: 3, 2, in: foo3)) |
| } |
| |
| SimpleMathTests.test("FunctionCall") { |
| func foo(_ x: Float, _ y: Float) -> Float { |
| return 3 * x + { $0 * 3 }(3) * y |
| } |
| expectEqual((3, 9), gradient(at: 3, 4, in: foo)) |
| expectEqual(3, gradient(at: 3) { x in foo(x, 4) }) |
| } |
| |
| SimpleMathTests.test("ResultSelection") { |
| func tuple(_ x: Float, _ y: Float) -> (Float, Float) { |
| return (x + 1, y + 2) |
| } |
| expectEqual((1, 0), gradient(at: 3, 3, in: { x, y in tuple(x, y).0 })) |
| expectEqual((0, 1), gradient(at: 3, 3, in: { x, y in tuple(x, y).1 })) |
| |
| func tupleGeneric<T>(_ x: T, _ y: T) -> (T, T) { |
| return (x, y) |
| } |
| func tupleGenericFirst<T>(_ x: T, _ y: T) -> T { tupleGeneric(x, y).0 } |
| func tupleGenericSecond<T>(_ x: T, _ y: T) -> T { tupleGeneric(x, y).1 } |
| expectEqual((1, 0), gradient(at: 3, 3, in: tupleGenericFirst)) |
| expectEqual((0, 1), gradient(at: 3, 3, in: tupleGenericSecond)) |
| } |
| |
| SimpleMathTests.test("MultipleResults") { |
| // Test function returning a tuple of active results. |
| func tuple(_ x: Float, _ y: Float) -> (Float, Float) { |
| return (x, y) |
| } |
| func multiply(_ x: Float, _ y: Float) -> Float { |
| let z = tuple(x, y) |
| // Note: both results (tuple elements) are active. |
| return z.0 * z.1 |
| } |
| expectEqual((4, 3), gradient(at: 3, 4, in: multiply)) |
| expectEqual((10, 5), gradient(at: 5, 10, in: multiply)) |
| |
| // Test function with multiple `inout` parameters. |
| func swap(_ x: inout Float, _ y: inout Float) { |
| let tmp = x; x = y; y = tmp |
| } |
| func multiply_swap(_ x: Float, _ y: Float) -> Float { |
| var tuple = (x, y) |
| swap(&tuple.0, &tuple.1) |
| return tuple.0 * tuple.1 |
| } |
| expectEqual((4, 3), gradient(at: 3, 4, in: multiply_swap)) |
| expectEqual((10, 5), gradient(at: 5, 10, in: multiply_swap)) |
| |
| // Test function with multiple `inout` parameters. |
| func swapGeneric<T>(_ x: inout T, _ y: inout T) { |
| let tmp = x; x = y; y = tmp |
| } |
| func multiply_swapGeneric(_ x: Float, _ y: Float) -> Float { |
| var tuple = (x, y) |
| swapGeneric(&tuple.0, &tuple.1) |
| return tuple.0 * tuple.1 |
| } |
| expectEqual((4, 3), gradient(at: 3, 4, in: multiply_swapGeneric)) |
| expectEqual((10, 5), gradient(at: 5, 10, in: multiply_swapGeneric)) |
| |
| // Test function with multiple `inout` parameters and a formal result. |
| func swapAndReturnProduct(_ x: inout Float, _ y: inout Float) -> Float { |
| let tmp = x |
| x = y |
| y = tmp |
| return x * y |
| } |
| func multiply_swapAndReturnProduct(_ x: Float, _ y: Float) -> Float { |
| var x2 = x |
| var y2 = y |
| let result = swapAndReturnProduct(&x2, &y2) |
| return result |
| } |
| expectEqual((4, 3), gradient(at: 3, 4, in: multiply_swapAndReturnProduct)) |
| expectEqual((4, 3), gradient(at: 3, 4, in: multiply_swapAndReturnProduct)) |
| } |
| |
| SimpleMathTests.test("CaptureLocal") { |
| let z: Float = 10 |
| func foo(_ x: Float) -> Float { |
| return z * x |
| } |
| expectEqual(10, gradient(at: 0, in: foo)) |
| } |
| |
| var globalVar: Float = 10 |
| SimpleMathTests.test("CaptureGlobal") { |
| func foo(x: Float) -> Float { |
| globalVar += 20 |
| return globalVar * x |
| } |
| expectEqual(30, gradient(at: 0, in: foo)) |
| } |
| |
| var foo_diffable: @differentiable (Float) -> (Float) |
| = differentiableFunction { x in (x * x, { v in 2 * x * v }) } |
| SimpleMathTests.test("GlobalDiffableFunc") { |
| expectEqual(2, gradient(at: 1, in: foo_diffable)) |
| expectEqual(2, gradient(at: 1, in: { x in foo_diffable(x) })) |
| expectEqual(1, gradient(at: 1, in: { (x: Float) -> Float in |
| foo_diffable = { x in x + 1 } |
| return foo_diffable(x) |
| })) |
| expectEqual(1, gradient(at: 1, in: foo_diffable)) |
| } |
| |
| SimpleMathTests.test("Mutation") { |
| func fourthPower(x: Float) -> Float { |
| var a = x |
| a = a * x |
| a = a * x |
| return a * x |
| } |
| expectEqual(4 * 27, gradient(at: 3, in: fourthPower)) |
| } |
| |
| SimpleMathTests.test("Tuple") { |
| // TF-945: Nested tuple projections. |
| func nested(_ x: Float) -> Float { |
| var tuple = (1, 1, ((x, 1), 1)) |
| return tuple.2.0.0 |
| } |
| expectEqual(1, gradient(at: 3, in: nested)) |
| } |
| |
| SimpleMathTests.test("TupleMutation") { |
| func foo(_ x: Float) -> Float { |
| var tuple = (x, x) |
| tuple.0 = tuple.0 * x |
| return x * tuple.0 |
| } |
| expectEqual(27, gradient(at: 3, in: foo)) |
| |
| func fifthPower(_ x: Float) -> Float { |
| var tuple = (x, x) |
| tuple.0 = tuple.0 * x |
| tuple.1 = tuple.0 * x |
| return tuple.0 * tuple.1 |
| } |
| expectEqual(405, gradient(at: 3, in: fifthPower)) |
| |
| func nested(_ x: Float) -> Float { |
| var tuple = ((x, x), x) |
| tuple.0.0 = tuple.0.0 * x |
| tuple.0.1 = tuple.0.0 * x |
| return tuple.0.0 * tuple.0.1 |
| } |
| expectEqual(405, gradient(at: 3, in: nested)) |
| |
| func generic<T: Differentiable & AdditiveArithmetic>(_ x: T) -> T { |
| var tuple = (x, x) |
| return tuple.0 |
| } |
| expectEqual(1, gradient(at: 3.0, in: generic)) |
| |
| // FIXME(TF-1033): Fix forward-mode ownership error for tuple with non-active |
| // initial values. |
| /* |
| func genericInitialNonactive<T: Differentiable & AdditiveArithmetic>( |
| _ x: T |
| ) -> T { |
| var tuple = (T.zero, T.zero) |
| tuple.0 = x |
| tuple.1 = x |
| return tuple.0 |
| } |
| expectEqual(1, gradient(at: 3.0, in: genericInitialNonactive)) |
| */ |
| } |
| |
| // Tests TF-321. |
| SimpleMathTests.test("TupleNonDifferentiableElements") { |
| // TF-964: Test tuple with non-tuple-typed adjoint value. |
| func tupleLet(_ x: Tracked<Float>) -> Tracked<Float> { |
| let tuple = (2 * x, 1) |
| return tuple.0 |
| } |
| expectEqual((8, 2), valueWithGradient(at: 4, in: tupleLet)) |
| |
| func tupleVar(_ x: Tracked<Float>) -> Tracked<Float> { |
| var tuple = (x, 1) |
| tuple.0 = x |
| tuple.1 = 1 |
| return tuple.0 |
| } |
| expectEqual((3, 1), valueWithGradient(at: 3, in: tupleVar)) |
| |
| func nested(_ x: Tracked<Float>) -> Tracked<Float> { |
| // Convoluted function computing `x * x`. |
| var tuple: (Int, (Int, Tracked<Float>), Tracked<Float>) = (1, (1, 0), 0) |
| tuple.0 = 1 |
| tuple.1.0 = 1 |
| tuple.1.1 = x |
| tuple.2 = x |
| return tuple.1.1 * tuple.2 |
| } |
| expectEqual((16, 8), valueWithGradient(at: 4, in: nested)) |
| |
| struct Wrapper<T> { |
| @differentiable(where T : Differentiable) |
| func baz(_ x: T) -> T { |
| var tuple = (1, 1, x, 1) |
| tuple.0 = 1 |
| tuple.2 = x |
| tuple.3 = 1 |
| return tuple.2 |
| } |
| } |
| func wrapper(_ x: Tracked<Float>) -> Tracked<Float> { |
| let w = Wrapper<Tracked<Float>>() |
| return w.baz(x) |
| } |
| expectEqual((3, 1), valueWithGradient(at: 3, in: wrapper)) |
| } |
| |
| // Tests TF-21. |
| SimpleMathTests.test("StructMemberwiseInitializer") { |
| struct Foo : AdditiveArithmetic, Differentiable { |
| var stored: Float |
| var computed: Float { |
| return stored * stored |
| } |
| } |
| |
| // Test direct `init` reference. |
| expectEqual(10, pullback(at: 4, in: Foo.init)(.init(stored: 10))) |
| |
| let 𝛁foo = pullback(at: Float(4), in: { input -> Foo in |
| let foo = Foo(stored: input) |
| let foo2 = foo + foo |
| return Foo(stored: foo2.stored) |
| })(Foo.TangentVector(stored: 1)) |
| expectEqual(2, 𝛁foo) |
| |
| let 𝛁computed = gradient(at: Float(4)) { input -> Float in |
| let foo = Foo(stored: input) |
| return foo.computed |
| } |
| expectEqual(8, 𝛁computed) |
| |
| let 𝛁product = gradient(at: Float(4)) { input -> Float in |
| let foo = Foo(stored: input) |
| return foo.computed * foo.stored |
| } |
| expectEqual(48, 𝛁product) |
| |
| struct Custom : AdditiveArithmetic, Differentiable { |
| var x: Float |
| |
| // Custom initializer with `@differentiable`. |
| @differentiable |
| init(x: Float) { |
| self.x = x |
| } |
| } |
| |
| let 𝛁custom = pullback(at: Float(4), in: { input -> Custom in |
| let foo = Custom(x: input) |
| return foo + foo |
| })(Custom.TangentVector(x: 1)) |
| expectEqual(2, 𝛁custom) |
| } |
| |
| // Tests TF-319: struct with non-differentiable constant stored property. |
| SimpleMathTests.test("StructConstantStoredProperty") { |
| struct TF_319 : Differentiable { |
| var x: Float |
| @noDerivative let constant = Float(2) |
| |
| @differentiable |
| init(x: Float) { |
| self.x = x |
| } |
| |
| @differentiable(wrt: (self, input)) |
| func applied(to input: Float) -> Float { |
| return x * constant * input |
| } |
| } |
| func testStructInit(to input: Float) -> Float { |
| let model = TF_319(x: 10) |
| return model.applied(to: input) |
| } |
| expectEqual(TF_319.TangentVector(x: 6), |
| gradient(at: TF_319(x: 10), in: { $0.applied(to: 3) })) |
| expectEqual(20, gradient(at: 3, in: testStructInit)) |
| } |
| |
| SimpleMathTests.test("StructMutation") { |
| struct Point : AdditiveArithmetic, Differentiable { |
| var x: Float |
| var y: Float |
| var z: Float |
| } |
| |
| func double(_ input: Float) -> Point { |
| let point = Point(x: input, y: input, z: input) |
| return point + point |
| } |
| expectEqual(6, pullback(at: 4, in: double)(Point(x: 1, y: 1, z: 1))) |
| |
| func fifthPower(_ input: Float) -> Float { |
| var point = Point(x: input, y: input, z: input) |
| point.x = point.x * input |
| point.y = point.x * input |
| return point.x * point.y |
| } |
| expectEqual(405, gradient(at: 3, in: fifthPower)) |
| |
| func mix(_ input: Float) -> Float { |
| var tuple = (point: Point(x: input, y: input, z: input), float: input) |
| tuple.point.x = tuple.point.x * tuple.float |
| tuple.point.y = tuple.point.x * input |
| return tuple.point.x * tuple.point.y |
| } |
| expectEqual(405, gradient(at: 3, in: mix)) |
| |
| // Test TF-282. |
| struct Add : Differentiable { |
| var bias: Float |
| func applied(to input: Float) -> Float { |
| var tmp = input |
| tmp = tmp + bias |
| return tmp |
| } |
| } |
| let model = Add(bias: 1) |
| expectEqual(Add.TangentVector(bias: 1), |
| gradient(at: model) { m in m.applied(to: 1) }) |
| } |
| |
| SimpleMathTests.test("StructGeneric") { |
| struct Generic<T : AdditiveArithmetic & Differentiable> : AdditiveArithmetic, Differentiable { |
| var x: T |
| var y: T |
| var z: T |
| } |
| |
| let 𝛁generic = pullback(at: Float(3), in: { input -> Generic<Float> in |
| var generic = Generic(x: input, y: input, z: input) |
| return generic |
| })(Generic<Float>.TangentVector(x: 1, y: 1, z: 1)) |
| expectEqual(3, 𝛁generic) |
| |
| func fifthPower(_ input: Float) -> Float { |
| var generic = Generic(x: input, y: input, z: input) |
| generic.x = generic.x * input |
| generic.y = generic.x * input |
| return generic.x * generic.y |
| } |
| expectEqual(405, gradient(at: 3, in: fifthPower)) |
| } |
| |
| SimpleMathTests.test("StructWithNoDerivativeProperty") { |
| struct NoDerivativeProperty : Differentiable { |
| var x: Float |
| @noDerivative var y: Float |
| } |
| expectEqual( |
| NoDerivativeProperty.TangentVector(x: 1), |
| gradient(at: NoDerivativeProperty(x: 1, y: 1)) { s -> Float in |
| var tmp = s |
| tmp.y = tmp.x |
| return tmp.x |
| } |
| ) |
| } |
| |
| SimpleMathTests.test("SubsetIndices") { |
| func grad(_ lossFunction: @differentiable (Float, Float) -> Float) -> Float { |
| return gradient(at: 1) { x in lossFunction(x * x, 10.0) } |
| } |
| expectEqual(2, grad { x, y in x + y }) |
| |
| func gradWRTNonDiff(_ lossFunction: @differentiable (Float, @noDerivative Int) -> Float) -> Float { |
| return gradient(at: 2) { x in lossFunction(x * x, 10) } |
| } |
| expectEqual(4, gradWRTNonDiff { x, y in x + Float(y) }) |
| } |
| |
| SimpleMathTests.test("ForceUnwrapping") { |
| func forceUnwrap<T: Differentiable & FloatingPoint>(_ t: T) |
| -> (T, Float) where T == T.TangentVector { |
| gradient(at: t, Float(1)) { (x, y) in |
| (x as! Float) * y |
| } |
| } |
| expectEqual((1, 2), forceUnwrap(Float(2))) |
| } |
| |
| SimpleMathTests.test("Adjoint value accumulation for aggregate lhs and concrete rhs") { |
| // TF-943: Test adjoint value accumulation for aggregate lhs and concrete rhs. |
| struct SmallTestModel : Differentiable { |
| public var stored: Float = 3.0 |
| @differentiable public func callAsFunction() -> Float { return stored } |
| } |
| |
| func doubled(_ model: SmallTestModel) -> Float{ |
| return model() + model.stored |
| } |
| let grads = gradient(at: SmallTestModel(), in: doubled) |
| expectEqual(2.0, grads.stored) |
| } |
| |
| // CHECK-LABEL: sil private [ossa] @AD__${{.*}}doubled{{.*}}pullback_src_0_wrt_0 : $@convention(thin) (Float, @owned {{.*}}) -> SmallTestModel.TangentVector { |
| // CHECK: bb0([[DX:%.*]] : $Float, [[PB_STRUCT:%.*]] : {{.*}}): |
| // CHECK: ([[PB0:%.*]], [[PB1:%.*]]) = destructure_struct [[PB_STRUCT]] |
| // CHECK: [[ADJ_TUPLE:%.*]] = apply [[PB1]]([[DX]]) : $@callee_guaranteed (Float) -> (Float, Float) |
| // CHECK: ([[TMP0:%.*]], [[ADJ_CONCRETE:%.*]]) = destructure_tuple [[ADJ_TUPLE]] : $(Float, Float) |
| // CHECK: [[TMP1:%.*]] = apply [[PB0]]([[TMP0]]) : $@callee_guaranteed (Float) -> SmallTestModel.TangentVector |
| // CHECK: [[ADJ_STRUCT_FIELD:%.*]] = destructure_struct [[TMP1]] : $SmallTestModel.TangentVector |
| // CHECK: [[TMP_RES:%.*]] = alloc_stack $Float |
| // CHECK: [[TMP_ADJ_STRUCT_FIELD:%.*]] = alloc_stack $Float |
| // CHECK: [[TMP_ADJ_CONCRETE:%.*]] = alloc_stack $Float |
| // CHECK: store [[ADJ_STRUCT_FIELD]] to [trivial] [[TMP_ADJ_STRUCT_FIELD]] : $*Float |
| // CHECK: store [[ADJ_CONCRETE]] to [trivial] [[TMP_ADJ_CONCRETE]] : $*Float |
| // CHECK: [[PLUS_EQUAL:%.*]] = witness_method $Float, #AdditiveArithmetic."+" |
| // CHECK: %{{.*}} = apply [[PLUS_EQUAL]]<Float>([[TMP_RES]], [[TMP_ADJ_CONCRETE]], [[TMP_ADJ_STRUCT_FIELD]], {{.*}}) |
| // CHECK: [[RES:%.*]] = load [trivial] [[TMP_RES]] : $*Float |
| // CHECK: [[RES_STRUCT:%.*]] = struct $SmallTestModel.TangentVector ([[RES]] : $Float) |
| // CHECK: return [[RES_STRUCT]] : $SmallTestModel.TangentVector |
| // CHECK: } |
| |
| runAllTests() |