From 87ded82acb5eec714e6c2fdae13b24699e2044ca Mon Sep 17 00:00:00 2001 From: Marc Rasi Date: Fri, 17 Apr 2020 16:48:48 -0700 Subject: [PATCH 1/3] [AutoDiff upstream] add more validation tests --- .../derivative_registration.swift | 196 +++++++++ .../differentiable_property.swift | 148 +++++++ .../validation-test/existential.swift | 26 ++ test/AutoDiff/validation-test/method.swift | 395 ++++++++++++++++++ .../validation-test/repeated_calls.swift | 19 + .../separate_tangent_type.swift | 50 +++ .../validation-test/superset_adjoint.swift | 82 ++++ 7 files changed, 916 insertions(+) create mode 100644 test/AutoDiff/validation-test/derivative_registration.swift create mode 100644 test/AutoDiff/validation-test/differentiable_property.swift create mode 100644 test/AutoDiff/validation-test/existential.swift create mode 100644 test/AutoDiff/validation-test/method.swift create mode 100644 test/AutoDiff/validation-test/repeated_calls.swift create mode 100644 test/AutoDiff/validation-test/separate_tangent_type.swift create mode 100644 test/AutoDiff/validation-test/superset_adjoint.swift diff --git a/test/AutoDiff/validation-test/derivative_registration.swift b/test/AutoDiff/validation-test/derivative_registration.swift new file mode 100644 index 0000000000000..81b761f754047 --- /dev/null +++ b/test/AutoDiff/validation-test/derivative_registration.swift @@ -0,0 +1,196 @@ +// RUN: %target-run-simple-swift +// REQUIRES: executable_test + +import StdlibUnittest +import DifferentiationUnittest + +var DerivativeRegistrationTests = TestSuite("DerivativeRegistration") + +@_semantics("autodiff.opaque") +func unary(x: Tracked) -> Tracked { + return x +} +@derivative(of: unary) +func _vjpUnary(x: Tracked) -> (value: Tracked, pullback: (Tracked) -> Tracked) { + return (value: x, pullback: { v in v }) +} +DerivativeRegistrationTests.testWithLeakChecking("UnaryFreeFunction") { + expectEqual(1, gradient(at: 3.0, in: unary)) +} + +@_semantics("autodiff.opaque") +func multiply(_ x: Tracked, _ y: Tracked) -> Tracked { + return x * y +} +@derivative(of: multiply) +func _vjpMultiply(_ x: Tracked, _ y: Tracked) + -> (value: Tracked, pullback: (Tracked) -> (Tracked, Tracked)) { + return (x * y, { v in (v * y, v * x) }) +} +DerivativeRegistrationTests.testWithLeakChecking("BinaryFreeFunction") { + expectEqual((3.0, 2.0), gradient(at: 2.0, 3.0, in: { x, y in multiply(x, y) })) +} + +struct Wrapper : Differentiable { + var float: Tracked +} + +extension Wrapper { + @_semantics("autodiff.opaque") + init(_ x: Tracked, _ y: Tracked) { + self.float = x * y + } + + @derivative(of: init(_:_:)) + static func _vjpInit(_ x: Tracked, _ y: Tracked) + -> (value: Self, pullback: (TangentVector) -> (Tracked, Tracked)) { + return (.init(x, y), { v in (v.float * y, v.float * x) }) + } +} +DerivativeRegistrationTests.testWithLeakChecking("Initializer") { + let v = Wrapper.TangentVector(float: 1) + let (𝛁x, 𝛁y) = pullback(at: 3, 4, in: { x, y in Wrapper(x, y) })(v) + expectEqual(4, 𝛁x) + expectEqual(3, 𝛁y) +} + +extension Wrapper { + @_semantics("autodiff.opaque") + static func multiply(_ x: Tracked, _ y: Tracked) -> Tracked { + return x * y + } + + @derivative(of: multiply) + static func _vjpMultiply(_ x: Tracked, _ y: Tracked) + -> (value: Tracked, pullback: (Tracked) -> (Tracked, Tracked)) { + return (x * y, { v in (v * y, v * x) }) + } +} +DerivativeRegistrationTests.testWithLeakChecking("StaticMethod") { + expectEqual((3.0, 2.0), gradient(at: 2.0, 3.0, in: { x, y in Wrapper.multiply(x, y) })) +} + +extension Wrapper { + @_semantics("autodiff.opaque") + func multiply(_ x: Tracked) -> Tracked { + return float * x + } + + @derivative(of: multiply) + func _vjpMultiply(_ x: Tracked) + -> (value: Tracked, pullback: (Tracked) -> (Wrapper.TangentVector, Tracked)) { + return (float * x, { v in + (TangentVector(float: v * x), v * self.float) + }) + } +} +DerivativeRegistrationTests.testWithLeakChecking("InstanceMethod") { + let x: Tracked = 2 + let wrapper = Wrapper(float: 3) + let (𝛁wrapper, 𝛁x) = gradient(at: wrapper, x) { wrapper, x in wrapper.multiply(x) } + expectEqual(Wrapper.TangentVector(float: 2), 𝛁wrapper) + expectEqual(3, 𝛁x) +} + +extension Wrapper { + subscript(_ x: Tracked) -> Tracked { + @_semantics("autodiff.opaque") + get { float * x } + set {} + } + + @derivative(of: subscript(_:)) + func _vjpSubscript(_ x: Tracked) + -> (value: Tracked, pullback: (Tracked) -> (Wrapper.TangentVector, Tracked)) { + return (self[x], { v in + (TangentVector(float: v * x), v * self.float) + }) + } +} +DerivativeRegistrationTests.testWithLeakChecking("Subscript") { + let x: Tracked = 2 + let wrapper = Wrapper(float: 3) + let (𝛁wrapper, 𝛁x) = gradient(at: wrapper, x) { wrapper, x in wrapper[x] } + expectEqual(Wrapper.TangentVector(float: 2), 𝛁wrapper) + expectEqual(3, 𝛁x) +} + +extension Wrapper { + var computedProperty: Tracked { + @_semantics("autodiff.opaque") + get { float * float } + set {} + } + + @derivative(of: computedProperty) + func _vjpComputedProperty() + -> (value: Tracked, pullback: (Tracked) -> Wrapper.TangentVector) { + return (computedProperty, { [f = self.float] v in + TangentVector(float: v * (f + f)) + }) + } +} +DerivativeRegistrationTests.testWithLeakChecking("ComputedProperty") { + let wrapper = Wrapper(float: 3) + let 𝛁wrapper = gradient(at: wrapper) { wrapper in wrapper.computedProperty } + expectEqual(Wrapper.TangentVector(float: 6), 𝛁wrapper) +} + +struct Generic { + @differentiable // derivative generic signature: none + func instanceMethod(_ x: Tracked) -> Tracked { + x + } +} +extension Generic { + @derivative(of: instanceMethod) // derivative generic signature: + func vjpInstanceMethod(_ x: Tracked) + -> (value: Tracked, pullback: (Tracked) -> Tracked) { + (x, { v in 1000 }) + } +} +DerivativeRegistrationTests.testWithLeakChecking("DerivativeGenericSignature") { + let generic = Generic() + let x: Tracked = 3 + let dx = gradient(at: x) { x in generic.instanceMethod(x) } + expectEqual(1000, dx) +} + +// When non-canonicalized generic signatures are used to compare derivative configurations, the +// `@differentiable` and `@derivative` attributes create separate derivatives, and we get a +// duplicate symbol error in TBDGen. +public protocol RefinesDifferentiable: Differentiable {} +extension Float: RefinesDifferentiable {} +@differentiable(where T: Differentiable, T: RefinesDifferentiable) +public func nonCanonicalizedGenSigComparison(_ t: T) -> T { t } +@derivative(of: nonCanonicalizedGenSigComparison) +public func dNonCanonicalizedGenSigComparison(_ t: T) + -> (value: T, pullback: (T.TangentVector) -> T.TangentVector) +{ + (t, { _ in T.TangentVector.zero }) +} +DerivativeRegistrationTests.testWithLeakChecking("NonCanonicalizedGenericSignatureComparison") { + let dx = gradient(at: Float(0), in: nonCanonicalizedGenSigComparison) + // Expect that we use the custom registered derivative, not a generated derivative (which would + // give a gradient of 1). + expectEqual(0, dx) +} + +// Test derivatives of default implementations. +protocol HasADefaultImplementation { + func req(_ x: Tracked) -> Tracked +} +extension HasADefaultImplementation { + func req(_ x: Tracked) -> Tracked { x } + @derivative(of: req) + func req(_ x: Tracked) -> (value: Tracked, pullback: (Tracked) -> Tracked) { + (x, { 10 * $0 }) + } +} +struct StructConformingToHasADefaultImplementation : HasADefaultImplementation {} +DerivativeRegistrationTests.testWithLeakChecking("DerivativeOfDefaultImplementation") { + let dx = gradient(at: Tracked(0)) { StructConformingToHasADefaultImplementation().req($0) } + expectEqual(Tracked(10), dx) +} + +runAllTests() diff --git a/test/AutoDiff/validation-test/differentiable_property.swift b/test/AutoDiff/validation-test/differentiable_property.swift new file mode 100644 index 0000000000000..9f9dfcf9ce4ec --- /dev/null +++ b/test/AutoDiff/validation-test/differentiable_property.swift @@ -0,0 +1,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 +} + +extension TangentSpace : Differentiable { + typealias TangentVector = TangentSpace +} + +struct Space { + /// `x` is a computed property with a custom vjp. + var x: Tracked { + @differentiable + get { storedX } + set { storedX = newValue } + } + + @derivative(of: x) + func vjpX() -> (value: Tracked, pullback: (Tracked) -> TangentSpace) { + return (x, { v in TangentSpace(x: v, y: 0) } ) + } + + private var storedX: Tracked + + @differentiable + var y: Tracked + + init(x: Tracked, y: Tracked) { + self.storedX = x + self.y = y + } +} + +extension Space : Differentiable { + typealias TangentVector = TangentSpace + mutating func move(along direction: TangentSpace) { + x.move(along: direction.x) + y.move(along: direction.y) + } +} + +E2EDifferentiablePropertyTests.testWithLeakChecking("computed property") { + let actualGrad = gradient(at: Space(x: 0, y: 0)) { (point: Space) -> Tracked 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 in + return 3 * point.y + } + let expectedGrad = TangentSpace(x: 0, y: 3) + expectEqual(expectedGrad, actualGrad) +} + +struct GenericMemberWrapper : Differentiable { + // Stored property. + @differentiable + 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>(x: 1)) { point in + return 2 * point.x + } + let expectedGrad = GenericMemberWrapper>.TangentVector(x: 2) + expectEqual(expectedGrad, actualGrad) +} + +struct ProductSpaceSelfTangent : AdditiveArithmetic { + let x, y: Tracked +} + +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 in + return 5 * point.y + } + let expectedGrad = ProductSpaceSelfTangent(x: 0, y: 5) + expectEqual(expectedGrad, actualGrad) +} + +struct ProductSpaceOtherTangentTangentSpace : AdditiveArithmetic { + let x, y: Tracked +} + +extension ProductSpaceOtherTangentTangentSpace : Differentiable { + typealias TangentVector = ProductSpaceOtherTangentTangentSpace +} + +struct ProductSpaceOtherTangent { + var x, y: Tracked +} + +extension ProductSpaceOtherTangent : Differentiable { + typealias TangentVector = ProductSpaceOtherTangentTangentSpace + mutating func move(along direction: ProductSpaceOtherTangentTangentSpace) { + x.move(along: direction.x) + y.move(along: direction.y) + } +} + +E2EDifferentiablePropertyTests.testWithLeakChecking("fieldwise product space, other tangent") { + let actualGrad = gradient( + at: ProductSpaceOtherTangent(x: 0, y: 0) + ) { (point: ProductSpaceOtherTangent) -> Tracked 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 + @differentiable + var computed: Tracked { + 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() diff --git a/test/AutoDiff/validation-test/existential.swift b/test/AutoDiff/validation-test/existential.swift new file mode 100644 index 0000000000000..fee835b9ff76d --- /dev/null +++ b/test/AutoDiff/validation-test/existential.swift @@ -0,0 +1,26 @@ +// RUN: %target-run-simple-swift +// REQUIRES: executable_test + +import StdlibUnittest +import DifferentiationUnittest + +var ExistentialTests = TestSuite("Existential") + +protocol A { + @differentiable(wrt: x) + func a(_ x: Tracked) -> Tracked +} +func b(g: A) -> Tracked { + return gradient(at: 3) { x in g.a(x) } +} + +struct B : A { + @differentiable(wrt: x) + func a(_ x: Tracked) -> Tracked { return x * 5 } +} + +ExistentialTests.testWithLeakChecking("Existential method VJP") { + expectEqual(5.0, b(g: B())) +} + +runAllTests() diff --git a/test/AutoDiff/validation-test/method.swift b/test/AutoDiff/validation-test/method.swift new file mode 100644 index 0000000000000..e5c2a5fa09b40 --- /dev/null +++ b/test/AutoDiff/validation-test/method.swift @@ -0,0 +1,395 @@ +// RUN: %target-run-simple-swift +// REQUIRES: executable_test + +import StdlibUnittest +import DifferentiationUnittest + +var MethodTests = TestSuite("Method") + +// ==== Tests with generated adjoint ==== + +struct Parameter : Equatable { + private let storedX: Tracked + @differentiable(wrt: (self)) + var x: Tracked { + return storedX + } + + init(x: Tracked) { + storedX = x + } + + @derivative(of: x) + func vjpX() -> (value: Tracked, pullback: (Tracked) -> Parameter) { + return (x, { dx in Parameter(x: dx) } ) + } + + @derivative(of: x) + func jvpX() -> (value: Tracked, differential: (Parameter) -> Tracked) { + return (x, { $0.x }) + } +} + +extension Parameter { + func squared() -> Tracked { + return x * x + } + + static func squared(p: Parameter) -> Tracked { + return p.x * p.x + } + + func multiplied(with other: Tracked) -> Tracked { + return x * other + } + + static func * (_ a: Parameter, _ b: Parameter) -> Tracked { + return a.x * b.x + } +} + +extension Parameter : Differentiable, AdditiveArithmetic { + typealias TangentVector = Parameter + typealias Scalar = Tracked + typealias Shape = () + init(repeating repeatedValue: Tracked, shape: ()) { + self.init(x: repeatedValue) + } + static func + (lhs: Parameter, rhs: Parameter) -> Parameter { + return Parameter(x: lhs.x + rhs.x) + } + static func - (lhs: Parameter, rhs: Parameter) -> Parameter { + return Parameter(x: lhs.x - rhs.x) + } + static func * (lhs: Scalar, rhs: Parameter) -> Parameter { + return Parameter(x: lhs * rhs.x) + } + static var zero: Parameter { return Parameter(x: 0) } +} + +MethodTests.testWithLeakChecking( + "instance method with generated adjoint, called from differentated func" +) { + func f(_ p: Parameter) -> Tracked { + return 100 * p.squared() + } + expectEqual(Parameter(x: 4 * 100), gradient(at: Parameter(x: 2), in: f)) + expectEqual(Parameter(x: 40 * 100), gradient(at: Parameter(x: 20), in: f)) +} + +MethodTests.testWithLeakChecking( + "instance method with generated adjoint, differentiated directly" +) { + // This is our current syntax for taking gradients of instance methods + // directly. If/when we develop nicer syntax for this, change this test. + func g(p: Parameter) -> Tracked { p.squared() } + expectEqual(Parameter(x: 4), gradient(at: Parameter(x: 2), in: g)) + expectEqual(Parameter(x: 40), gradient(at: Parameter(x: 20), in: g)) +} + +MethodTests.testWithLeakChecking("instance method with generated adjoint, wrt only self") { + func f(_ p: Parameter) -> Tracked { + return 100 * p.multiplied(with: 200) + } + expectEqual(Parameter(x: 100 * 200), gradient(at: Parameter(x: 1), in: f)) + expectEqual(Parameter(x: 100 * 200), gradient(at: Parameter(x: 2), in: f)) +} + +MethodTests.testWithLeakChecking("instance method with generated adjoint, wrt only non-self") { + func f(_ other: Tracked) -> Tracked { + return 100 * Parameter(x: 200).multiplied(with: other) + } + expectEqual(100 * 200, gradient(at: 1, in: f)) + expectEqual(100 * 200, gradient(at: 2, in: f)) +} + +// FIXME: Add a binary differential operator. +// +// MethodTests.testWithLeakChecking( +// "instance method with generated adjoint, wrt self and non-self" +// ) { +// let g = #gradient({ (p: Parameter, o: Tracked) in p.multiplied(with: o) }) +// expectEqual((Parameter(x: 100), 200), g(Parameter(x: 200), 100)) +// expectEqual((Parameter(x: 200), 100), g(Parameter(x: 100), 200)) +// } + +MethodTests.testWithLeakChecking( + "static method with generated adjoint, called from differentiated func" +) { + func f(_ p: Parameter) -> Tracked { + return 100 * Parameter.squared(p: p) + } + expectEqual(Parameter(x: 4 * 100), gradient(at: Parameter(x: 2), in: f)) + expectEqual(Parameter(x: 40 * 100), gradient(at: Parameter(x: 20), in: f)) +} + +// TODO(SR-8699): Fix this test. +// MethodTests.testWithLeakChecking( +// "static method with generated adjoint, differentiated directly" +// ) { +// let grad = #gradient(Parameter.squared(p:)) +// expectEqual(Parameter(x: 4), grad(Parameter(x: 2))) +// expectEqual(Parameter(x: 40), grad(Parameter(x: 20))) +// } + +MethodTests.testWithLeakChecking("static method with generated adjoint, wrt only first param") { + func f(_ p: Parameter) -> Tracked { + return 100 * (p * Parameter(x: 200)) + } + expectEqual(Parameter(x: 100 * 200), gradient(at: Parameter(x: 1), in: f)) + expectEqual(Parameter(x: 100 * 200), gradient(at: Parameter(x: 2), in: f)) +} + +MethodTests.testWithLeakChecking("static method with generated adjoint, wrt only second param") { + func f(_ p: Parameter) -> Tracked { + return 100 * (Parameter(x: 200) * p) + } + expectEqual(Parameter(x: 100 * 200), gradient(at: Parameter(x: 1), in: f)) + expectEqual(Parameter(x: 100 * 200), gradient(at: Parameter(x: 2), in: f)) +} + +MethodTests.testWithLeakChecking("static method with generated adjoint, wrt all params") { + func g(a: Parameter, b: Parameter) -> Tracked { a * b } + expectEqual((Parameter(x: 100), Parameter(x: 200)), + gradient(at: Parameter(x: 200), Parameter(x: 100), in: g)) + expectEqual((Parameter(x: 200), Parameter(x: 100)), + gradient(at: Parameter(x: 100), Parameter(x: 200), in: g)) +} + +// ==== Tests with custom adjoint ==== + +// Test self-reordering thunk for jvp/vjp methods. +struct DiffWrtSelf : Differentiable { + @differentiable(wrt: (self, x, y)) + func call(_ x: T, _ y: U) -> T { + return x + } + @derivative(of: call) + func _jvpCall(_ x: T, _ y: U) + -> (value: T, differential: (DiffWrtSelf.TangentVector, T.TangentVector, U.TangentVector) -> T.TangentVector) { + return (x, { (dself, dx, dy) in dx }) + } + @derivative(of: call) + func _vjpCall(_ x: T, _ y: U) + -> (value: T, pullback: (T.TangentVector) -> (DiffWrtSelf.TangentVector, T.TangentVector, U.TangentVector)) { + return (x, { (.zero, $0, .zero) }) + } +} + +struct CustomParameter : Equatable { + let storedX: Tracked + @differentiable(wrt: (self)) + var x: Tracked { + return storedX + } + + init(x: Tracked) { + storedX = x + } + + @derivative(of: x) + func vjpX() -> (value: Tracked, pullback: (Tracked) -> CustomParameter) { + return (x, { dx in CustomParameter(x: dx) }) + } +} + +extension CustomParameter : Differentiable, AdditiveArithmetic { + typealias TangentVector = CustomParameter + typealias Scalar = Tracked + typealias Shape = () + init(repeating repeatedValue: Tracked, shape: ()) { + self.init(x: repeatedValue) + } + static func + (lhs: CustomParameter, rhs: CustomParameter) -> CustomParameter { + return CustomParameter(x: lhs.x + rhs.x) + } + static func - (lhs: CustomParameter, rhs: CustomParameter) -> CustomParameter { + return CustomParameter(x: lhs.x - rhs.x) + } + static func * (lhs: Scalar, rhs: CustomParameter) -> CustomParameter { + return CustomParameter(x: lhs * rhs.x) + } + static var zero: CustomParameter { return CustomParameter(x: 0) } +} + +extension Tracked where T : FloatingPoint { + func clamped(to limits: ClosedRange>) -> Tracked { + return min(max(self, limits.lowerBound), limits.upperBound) + } +} + +extension CustomParameter { + @differentiable(wrt: (self)) + func squared() -> Tracked { + return x * x + } + + @derivative(of: squared) + func dSquared() -> (value: Tracked, pullback: (Tracked) -> CustomParameter) { + return (squared(), { [x] v in CustomParameter(x: (2 * x).clamped(to: -10.0...10.0) * v) }) + } + + @differentiable + static func squared(p: CustomParameter) -> Tracked { + return p.x * p.x + } + + @derivative(of: squared) + static func dSquared( + _ p: CustomParameter + ) -> (value: Tracked, pullback: (Tracked) -> CustomParameter) { + return (p.x * p.x, { v in CustomParameter(x: (2 * p.x).clamped(to: -10.0...10.0) * v) }) + } + + // There is currently no way to define multiple custom VJPs wrt different + // parameters on the same func, so we define a copy of this func per adjoint. + + @differentiable(wrt: (self, other)) + func multiplied(with other: Tracked) -> Tracked { + return x * other + } + + @differentiable(wrt: (other)) + func multiplied_constSelf(with other: Tracked) -> Tracked { + return x * other + } + + @differentiable(wrt: (self)) + func multiplied_constOther(with other: Tracked) -> Tracked { + return x * other + } + + @derivative(of: multiplied) + func dMultiplied_wrtAll( + with other: Tracked + ) -> (value: Tracked, pullback: (Tracked) -> (CustomParameter, Tracked)) { + return (multiplied(with: other), + { [x] v in (CustomParameter(x: other.clamped(to: -10.0...10.0) * v), + x.clamped(to: -10.0...10.0) * v) }) + } + + @derivative(of: multiplied_constSelf, wrt: other) + func dMultiplied_wrtOther( + with other: Tracked + ) -> (value: Tracked, pullback: (Tracked) -> Tracked) { + let (r, pb) = dMultiplied_wrtAll(with: other) + return (r, { v in pb(v).1 }) + } + + @derivative(of: multiplied_constOther, wrt: self) + func dMultiplied_wrtSelf( + with other: Tracked + ) -> (value: Tracked, pullback: (Tracked) -> CustomParameter) { + let (r, pb) = dMultiplied_wrtAll(with: other) + return (r, { v in pb(v).0 }) + } + + @differentiable + static func multiply(_ lhs: CustomParameter, _ rhs: CustomParameter) + -> Tracked { + return lhs.x * rhs.x + } + + @differentiable(wrt: (rhs)) + static func multiply_constLhs(_ lhs: CustomParameter, _ rhs: CustomParameter) -> Tracked { + return lhs.x * rhs.x + } + + @derivative(of: multiply) + static func dMultiply_wrtAll(_ lhs: CustomParameter,_ rhs: CustomParameter) + -> (value: Tracked, pullback: (Tracked) -> (CustomParameter, CustomParameter)) { + let result = multiply(lhs, rhs) + return (result, { v in (CustomParameter(x: rhs.x.clamped(to: -10.0...10.0) * v), + CustomParameter(x: lhs.x.clamped(to: -10.0...10.0) * v)) }) + } + + @derivative(of: multiply_constLhs, wrt: rhs) + static func dMultiply_wrtRhs(_ lhs: CustomParameter, _ rhs: CustomParameter) + -> (value: Tracked, pullback: (Tracked) -> CustomParameter) { + let (r, pb) = dMultiply_wrtAll(lhs, rhs) + return (r, { v in pb(v).1 }) + } +} + +MethodTests.testWithLeakChecking( + "instance method with custom adjoint, called from differentated func" +) { + func f(_ p: CustomParameter) -> Tracked { + return 100 * p.squared() + } + expectEqual(CustomParameter(x: 4 * 100), gradient(at: CustomParameter(x: 2), in: f)) + expectEqual(CustomParameter(x: 10 * 100), gradient(at: CustomParameter(x: 20), in: f)) +} + +MethodTests.testWithLeakChecking("instance method with generated adjoint, differentated directly") { + // This is our current syntax for taking gradients of instance methods + // directly. If/when we develop nicer syntax for this, change this test. + func g(p: CustomParameter) -> Tracked { p.squared() } + expectEqual(CustomParameter(x: 4), gradient(at: CustomParameter(x: 2), in: g)) + expectEqual(CustomParameter(x: 10), gradient(at: CustomParameter(x: 20), in: g)) +} + +MethodTests.testWithLeakChecking("static method with custom adjoint, called from differentated func") { + func f(_ p: CustomParameter) -> Tracked { + return 100 * CustomParameter.squared(p: p) + } + expectEqual(CustomParameter(x: 4 * 100), gradient(at: CustomParameter(x: 2), in: f)) + expectEqual(CustomParameter(x: 10 * 100), gradient(at: CustomParameter(x: 20), in: f)) +} + +// TODO(SR-8699): Fix this test. +// MethodTests.testWithLeakChecking("static method with custom adjoint, differentiated directly") { +// let grad = #gradient(CustomParameter.squared(p:)) +// expectEqual(CustomParameter(x: 4), grad(CustomParameter(x: 2))) +// expectEqual(CustomParameter(x: 10), grad(CustomParameter(x: 20))) +// } + +MethodTests.testWithLeakChecking("instance method with custom adjoint, wrt only self") { + func f(_ p: CustomParameter) -> Tracked { + return 100 * p.multiplied_constOther(with: 200) + } + expectEqual(CustomParameter(x: 100 * 10), gradient(at: CustomParameter(x: 1), in: f)) + expectEqual(CustomParameter(x: 100 * 10), gradient(at: CustomParameter(x: 2), in: f)) +} + +MethodTests.testWithLeakChecking("instance method with custom adjoint, wrt only non-self") { + func f(_ other: Tracked) -> Tracked { + return 100 * CustomParameter(x: 200).multiplied_constSelf(with: other) + } + expectEqual(100 * 10, gradient(at: 1, in: f)) + expectEqual(100 * 10, gradient(at: 2, in: f)) +} + +MethodTests.testWithLeakChecking("instance method with custom adjoint, wrt self and non-self") { + func g(p: CustomParameter, o: Tracked) -> Tracked { p.multiplied(with: o) } + expectEqual((CustomParameter(x: 5), 10), gradient(at: CustomParameter(x: 100), 5, in: g)) + expectEqual((CustomParameter(x: 10), 5), gradient(at: CustomParameter(x: 5), 100, in: g)) +} + +MethodTests.testWithLeakChecking("static method with custom adjoint, wrt only lhs") { + func f(_ p: CustomParameter) -> Tracked { + return 100 * CustomParameter.multiply_constLhs(CustomParameter(x: 200), p) + } + expectEqual(CustomParameter(x: 100 * 10), gradient(at: CustomParameter(x: 1), in: f)) + expectEqual(CustomParameter(x: 100 * 10), gradient(at: CustomParameter(x: 2), in: f)) +} + +MethodTests.testWithLeakChecking("static method with custom adjoint, wrt only rhs") { + func f(_ p: CustomParameter) -> Tracked { + return 100 * CustomParameter.multiply_constLhs(CustomParameter(x: 200), p) + } + expectEqual(CustomParameter(x: 100 * 10), gradient(at: CustomParameter(x: 1), in: f)) + expectEqual(CustomParameter(x: 100 * 10), gradient(at: CustomParameter(x: 2), in: f)) +} + +MethodTests.testWithLeakChecking("static method with custom adjoint, wrt all") { + func f(_ a: CustomParameter, _ b: CustomParameter) -> Tracked { + return CustomParameter.multiply(a, b) + } + expectEqual((CustomParameter(x: 5), CustomParameter(x: 10)), + gradient(at: CustomParameter(x: 100), CustomParameter(x: 5), in: f)) + expectEqual((CustomParameter(x: 10), CustomParameter(x: 5)), + gradient(at: CustomParameter(x: 5), CustomParameter(x: 100), in: f)) +} + +runAllTests() diff --git a/test/AutoDiff/validation-test/repeated_calls.swift b/test/AutoDiff/validation-test/repeated_calls.swift new file mode 100644 index 0000000000000..ac1a495895e5e --- /dev/null +++ b/test/AutoDiff/validation-test/repeated_calls.swift @@ -0,0 +1,19 @@ +// RUN: %target-run-simple-swift +// REQUIRES: executable_test + +import StdlibUnittest +import DifferentiationUnittest + +var RepeatedCallsTests = TestSuite("RepeatedCalls") + +RepeatedCallsTests.testWithLeakChecking("Repeat") { + func mul2(_ x: Tracked) -> Tracked { + return 2 * x + } + func mul4(_ x: Tracked) -> Tracked { + return mul2(mul2(x)) + } + expectEqual(4, gradient(at: 0, in: mul4)) +} + +runAllTests() diff --git a/test/AutoDiff/validation-test/separate_tangent_type.swift b/test/AutoDiff/validation-test/separate_tangent_type.swift new file mode 100644 index 0000000000000..f5394d6c33379 --- /dev/null +++ b/test/AutoDiff/validation-test/separate_tangent_type.swift @@ -0,0 +1,50 @@ +// RUN: %target-run-simple-swift +// REQUIRES: executable_test + +import StdlibUnittest +#if os(macOS) +import Darwin.C +#else +import Glibc +#endif +import DifferentiationUnittest + +var SeparateTangentTypeTests = TestSuite("SeparateTangentType") + +struct DifferentiableSubset : Differentiable { + @differentiable(wrt: self) + var w: Tracked + @differentiable(wrt: self) + var b: Tracked + @noDerivative var flag: Bool + + struct TangentVector : Differentiable, AdditiveArithmetic { + typealias TangentVector = DifferentiableSubset.TangentVector + var w: Tracked + var b: Tracked + } + mutating func move(along v: TangentVector) { + w.move(along: v.w) + b.move(along: v.b) + } +} + +SeparateTangentTypeTests.testWithLeakChecking("Trivial") { + let x = DifferentiableSubset(w: 0, b: 1, flag: false) + let pb = pullback(at: x) { x in x } + expectEqual(pb(DifferentiableSubset.TangentVector.zero), DifferentiableSubset.TangentVector.zero) +} + +SeparateTangentTypeTests.testWithLeakChecking("Initialization") { + let x = DifferentiableSubset(w: 0, b: 1, flag: false) + let pb = pullback(at: x) { x in DifferentiableSubset(w: 1, b: 2, flag: true) } + expectEqual(pb(DifferentiableSubset.TangentVector.zero), DifferentiableSubset.TangentVector.zero) +} + +SeparateTangentTypeTests.testWithLeakChecking("SomeArithmetics") { + let x = DifferentiableSubset(w: 0, b: 1, flag: false) + let pb = pullback(at: x) { x in DifferentiableSubset(w: x.w * x.w, b: x.b * x.b, flag: true) } + expectEqual(pb(DifferentiableSubset.TangentVector.zero), DifferentiableSubset.TangentVector.zero) +} + +runAllTests() diff --git a/test/AutoDiff/validation-test/superset_adjoint.swift b/test/AutoDiff/validation-test/superset_adjoint.swift new file mode 100644 index 0000000000000..0bec03b6336fa --- /dev/null +++ b/test/AutoDiff/validation-test/superset_adjoint.swift @@ -0,0 +1,82 @@ +// RUN: %target-run-simple-swift +// REQUIRES: executable_test + +import StdlibUnittest +import DifferentiationUnittest + +var SupersetVJPTests = TestSuite("SupersetVJP") + +@differentiable(wrt: (x, y)) +func mulxy(_ x: Tracked, _ y: Tracked) -> Tracked { + // use control flow to prevent AD; NB fix when control flow is supported + if x > 1000 { + return y + } + return x * y +} +@derivative(of: mulxy) +func dmulxy( + _ x: Tracked, + _ y: Tracked +) -> (value: Tracked, pullback: (Tracked) -> (Tracked, Tracked)) { + return (mulxy(x, y), { v in (y * v, x * v) }) +} + +func calls_mulxy(_ x: Tracked, _ y: Tracked) -> Tracked { + return mulxy(x, y) +} + +SupersetVJPTests.testWithLeakChecking("Superset") { + expectEqual(3, gradient(at: 2) { x in mulxy(x, 3) }) +} + +SupersetVJPTests.testWithLeakChecking("SupersetNested") { + expectEqual(2, gradient(at: 3) { y in calls_mulxy(2, y) }) +} + +SupersetVJPTests.testWithLeakChecking("CrossModuleClosure") { + expectEqual(1, gradient(at: Tracked(1)) { x in x + 2 }) +} + +SupersetVJPTests.testWithLeakChecking("SubsetOfSubset") { + @differentiable(wrt: (x, z)) + func foo(_ x: Tracked, _ y: Tracked, _ z: Tracked) -> Tracked { + withoutDerivative(at: 0) + } + expectEqual(0, gradient(at: 0, in: { x in foo(x, 0, 0) })) +} + +SupersetVJPTests.test("ApplySubset") { + // TF-914 + @differentiable(wrt: x) + func foo(_ x: T, _ y: T, apply: @differentiable (T, T) -> T) -> T { + return apply(x, y) + } + expectEqual(1, gradient(at: Tracked(0)) { x in foo(x, 0) { $0 + $1 } }) +} + +// FIXME: The expression `(+) as @differentiable (Float, @noDerivative Float) -> Float)` +// forms a curry thunk of `Float.+` before conversion to @differentiable, and AD +// doesn't know how to differentiate the curry thunk, so it produces a +// "function is not differentiable" error. +// SupersetVJPTests.test("CrossModule") { +// let grad = gradient(at: Float(1), Float(2), in: (+) as @differentiable (Float, @noDerivative Float) -> Float) +// expectEqual(Float(1), grad) +// } + +@differentiable(wrt: (x, y)) +func x_T(_ x: Tracked, _ y: T) -> Tracked { + if x > 1000 { return x } + return x +} +@derivative(of: x_T) +func dx_T( + _ x: Tracked, _ y: T +) -> (value: Tracked, pullback: (Tracked) -> (Tracked, T.TangentVector)) { + return (x_T(x, y), { v in (x * v, .zero) }) +} +SupersetVJPTests.testWithLeakChecking("IndirectResults") { + expectEqual(2, gradient(at: 2) { x in x_T(x, Tracked(3)) }) +} + +runAllTests() From 9e5b4c107db6c06546fec36282235eebe5d1c8d5 Mon Sep 17 00:00:00 2001 From: Marc Rasi Date: Fri, 17 Apr 2020 20:51:04 -0700 Subject: [PATCH 2/3] fix test on iphonesimulator --- test/AutoDiff/validation-test/separate_tangent_type.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/AutoDiff/validation-test/separate_tangent_type.swift b/test/AutoDiff/validation-test/separate_tangent_type.swift index f5394d6c33379..30d77262162cb 100644 --- a/test/AutoDiff/validation-test/separate_tangent_type.swift +++ b/test/AutoDiff/validation-test/separate_tangent_type.swift @@ -2,7 +2,7 @@ // REQUIRES: executable_test import StdlibUnittest -#if os(macOS) +#if os(macOS) || os(iOS) || os(watchOS) || os(tvOS) import Darwin.C #else import Glibc From 2b0beec8088d14eb405107bf0dd76e08dd38373e Mon Sep 17 00:00:00 2001 From: Marc Rasi Date: Mon, 20 Apr 2020 13:34:13 -0700 Subject: [PATCH 3/3] enable tests --- test/AutoDiff/validation-test/method.swift | 43 +++++++++---------- .../validation-test/superset_adjoint.swift | 12 ++---- 2 files changed, 24 insertions(+), 31 deletions(-) diff --git a/test/AutoDiff/validation-test/method.swift b/test/AutoDiff/validation-test/method.swift index e5c2a5fa09b40..1f37af8230e68 100644 --- a/test/AutoDiff/validation-test/method.swift +++ b/test/AutoDiff/validation-test/method.swift @@ -103,15 +103,14 @@ MethodTests.testWithLeakChecking("instance method with generated adjoint, wrt on expectEqual(100 * 200, gradient(at: 2, in: f)) } -// FIXME: Add a binary differential operator. -// -// MethodTests.testWithLeakChecking( -// "instance method with generated adjoint, wrt self and non-self" -// ) { -// let g = #gradient({ (p: Parameter, o: Tracked) in p.multiplied(with: o) }) -// expectEqual((Parameter(x: 100), 200), g(Parameter(x: 200), 100)) -// expectEqual((Parameter(x: 200), 100), g(Parameter(x: 100), 200)) -// } +MethodTests.testWithLeakChecking( + "instance method with generated adjoint, wrt self and non-self" +) { + expectEqual( + (Parameter(x: 100), 200), gradient(at: Parameter(x: 200), 100) { $0.multiplied(with: $1) }) + expectEqual( + (Parameter(x: 200), 100), gradient(at: Parameter(x: 100), 200) { $0.multiplied(with: $1) }) +} MethodTests.testWithLeakChecking( "static method with generated adjoint, called from differentiated func" @@ -123,14 +122,12 @@ MethodTests.testWithLeakChecking( expectEqual(Parameter(x: 40 * 100), gradient(at: Parameter(x: 20), in: f)) } -// TODO(SR-8699): Fix this test. -// MethodTests.testWithLeakChecking( -// "static method with generated adjoint, differentiated directly" -// ) { -// let grad = #gradient(Parameter.squared(p:)) -// expectEqual(Parameter(x: 4), grad(Parameter(x: 2))) -// expectEqual(Parameter(x: 40), grad(Parameter(x: 20))) -// } +MethodTests.testWithLeakChecking( + "static method with generated adjoint, differentiated directly" +) { + expectEqual(Parameter(x: 4), gradient(at: Parameter(x: 2), in: Parameter.squared)) + expectEqual(Parameter(x: 40), gradient(at: Parameter(x: 20), in: Parameter.squared)) +} MethodTests.testWithLeakChecking("static method with generated adjoint, wrt only first param") { func f(_ p: Parameter) -> Tracked { @@ -337,12 +334,12 @@ MethodTests.testWithLeakChecking("static method with custom adjoint, called from expectEqual(CustomParameter(x: 10 * 100), gradient(at: CustomParameter(x: 20), in: f)) } -// TODO(SR-8699): Fix this test. -// MethodTests.testWithLeakChecking("static method with custom adjoint, differentiated directly") { -// let grad = #gradient(CustomParameter.squared(p:)) -// expectEqual(CustomParameter(x: 4), grad(CustomParameter(x: 2))) -// expectEqual(CustomParameter(x: 10), grad(CustomParameter(x: 20))) -// } +MethodTests.testWithLeakChecking("static method with custom adjoint, differentiated directly") { + expectEqual( + CustomParameter(x: 4), gradient(at: CustomParameter(x: 2), in: CustomParameter.squared)) + expectEqual( + CustomParameter(x: 10), gradient(at: CustomParameter(x: 20), in: CustomParameter.squared)) +} MethodTests.testWithLeakChecking("instance method with custom adjoint, wrt only self") { func f(_ p: CustomParameter) -> Tracked { diff --git a/test/AutoDiff/validation-test/superset_adjoint.swift b/test/AutoDiff/validation-test/superset_adjoint.swift index 0bec03b6336fa..bd250c367eda6 100644 --- a/test/AutoDiff/validation-test/superset_adjoint.swift +++ b/test/AutoDiff/validation-test/superset_adjoint.swift @@ -55,14 +55,10 @@ SupersetVJPTests.test("ApplySubset") { expectEqual(1, gradient(at: Tracked(0)) { x in foo(x, 0) { $0 + $1 } }) } -// FIXME: The expression `(+) as @differentiable (Float, @noDerivative Float) -> Float)` -// forms a curry thunk of `Float.+` before conversion to @differentiable, and AD -// doesn't know how to differentiate the curry thunk, so it produces a -// "function is not differentiable" error. -// SupersetVJPTests.test("CrossModule") { -// let grad = gradient(at: Float(1), Float(2), in: (+) as @differentiable (Float, @noDerivative Float) -> Float) -// expectEqual(Float(1), grad) -// } +SupersetVJPTests.test("CrossModule") { + let grad = gradient(at: Float(1)) { $0 + 2 } + expectEqual(Float(1), grad) +} @differentiable(wrt: (x, y)) func x_T(_ x: Tracked, _ y: T) -> Tracked {