[AutoDiff upstream] add more validation tests

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<Float>) -> Tracked<Float> {
+  return x
+}
+@derivative(of: unary)
+func _vjpUnary(x: Tracked<Float>) -> (value: Tracked<Float>, pullback: (Tracked<Float>) -> Tracked<Float>) {
+  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<Float>, _ y: Tracked<Float>) -> Tracked<Float> {
+  return x * y
+}
+@derivative(of: multiply)
+func _vjpMultiply(_ x: Tracked<Float>, _ y: Tracked<Float>)
+  -> (value: Tracked<Float>, pullback: (Tracked<Float>) -> (Tracked<Float>, Tracked<Float>)) {
+  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<Float>
+}
+
+extension Wrapper {
+  @_semantics("autodiff.opaque")
+  init(_ x: Tracked<Float>, _ y: Tracked<Float>) {
+    self.float = x * y
+  }
+
+  @derivative(of: init(_:_:))
+  static func _vjpInit(_ x: Tracked<Float>, _ y: Tracked<Float>)
+    -> (value: Self, pullback: (TangentVector) -> (Tracked<Float>, Tracked<Float>)) {
+    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<Float>, _ y: Tracked<Float>) -> Tracked<Float> {
+    return x * y
+  }
+
+  @derivative(of: multiply)
+  static func _vjpMultiply(_ x: Tracked<Float>, _ y: Tracked<Float>)
+    -> (value: Tracked<Float>, pullback: (Tracked<Float>) -> (Tracked<Float>, Tracked<Float>)) {
+    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<Float>) -> Tracked<Float> {
+    return float * x
+  }
+
+  @derivative(of: multiply)
+  func _vjpMultiply(_ x: Tracked<Float>)
+    -> (value: Tracked<Float>, pullback: (Tracked<Float>) -> (Wrapper.TangentVector, Tracked<Float>)) {
+    return (float * x, { v in
+      (TangentVector(float: v * x), v * self.float)
+    })
+  }
+}
+DerivativeRegistrationTests.testWithLeakChecking("InstanceMethod") {
+  let x: Tracked<Float> = 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<Float>) -> Tracked<Float> {
+    @_semantics("autodiff.opaque")
+    get { float * x }
+    set {}
+  }
+
+  @derivative(of: subscript(_:))
+  func _vjpSubscript(_ x: Tracked<Float>)
+    -> (value: Tracked<Float>, pullback: (Tracked<Float>) -> (Wrapper.TangentVector, Tracked<Float>)) {
+    return (self[x], { v in
+      (TangentVector(float: v * x), v * self.float)
+    })
+  }
+}
+DerivativeRegistrationTests.testWithLeakChecking("Subscript") {
+  let x: Tracked<Float> = 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<Float> {
+    @_semantics("autodiff.opaque")
+    get { float * float }
+    set {}
+  }
+
+  @derivative(of: computedProperty)
+  func _vjpComputedProperty()
+    -> (value: Tracked<Float>, pullback: (Tracked<Float>) -> 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<T> {
+  @differentiable // derivative generic signature: none
+  func instanceMethod(_ x: Tracked<Float>) -> Tracked<Float> {
+    x
+  }
+}
+extension Generic {
+  @derivative(of: instanceMethod) // derivative generic signature: <T>
+  func vjpInstanceMethod(_ x: Tracked<Float>)
+    -> (value: Tracked<Float>, pullback: (Tracked<Float>) -> Tracked<Float>) {
+    (x, { v in 1000 })
+  }
+}
+DerivativeRegistrationTests.testWithLeakChecking("DerivativeGenericSignature") {
+  let generic = Generic<Float>()
+  let x: Tracked<Float> = 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 { t }
+@derivative(of: nonCanonicalizedGenSigComparison)
+public func dNonCanonicalizedGenSigComparison<T: RefinesDifferentiable>(_ 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<Float>) -> Tracked<Float>
+}
+extension HasADefaultImplementation {
+  func req(_ x: Tracked<Float>) -> Tracked<Float> { x }
+  @derivative(of: req)
+  func req(_ x: Tracked<Float>) -> (value: Tracked<Float>, pullback: (Tracked<Float>) -> Tracked<Float>) {
+    (x, { 10 * $0 })
+  }
+}
+struct StructConformingToHasADefaultImplementation : HasADefaultImplementation {}
+DerivativeRegistrationTests.testWithLeakChecking("DerivativeOfDefaultImplementation") {
+  let dx = gradient(at: Tracked<Float>(0)) { StructConformingToHasADefaultImplementation().req($0) }
+  expectEqual(Tracked<Float>(10), dx)
+}
+
+runAllTests()

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<Float>
+}
+
+extension TangentSpace : Differentiable {
+  typealias TangentVector = TangentSpace
+}
+
+struct Space {
+  /// `x` is a computed property with a custom vjp.
+  var x: Tracked<Float> {
+    @differentiable
+    get { storedX }
+    set { storedX = newValue }
+  }
+
+  @derivative(of: x)
+  func vjpX() -> (value: Tracked<Float>, pullback: (Tracked<Float>) -> TangentSpace) {
+    return (x, { v in TangentSpace(x: v, y: 0) } )
+  }
+
+  private var storedX: Tracked<Float>
+
+  @differentiable
+  var y: Tracked<Float>
+
+  init(x: Tracked<Float>, y: Tracked<Float>) {
+    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<Float> 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<Float> in
+    return 3 * point.y
+  }
+  let expectedGrad = TangentSpace(x: 0, y: 3)
+  expectEqual(expectedGrad, actualGrad)
+}
+
+struct GenericMemberWrapper<T : Differentiable> : 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<Tracked<Float>>(x: 1)) { point in
+    return 2 * point.x
+  }
+  let expectedGrad = GenericMemberWrapper<Tracked<Float>>.TangentVector(x: 2)
+  expectEqual(expectedGrad, actualGrad)
+}
+
+struct ProductSpaceSelfTangent : AdditiveArithmetic {
+  let x, y: Tracked<Float>
+}
+
+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<Float> in
+    return 5 * point.y
+  }
+  let expectedGrad = ProductSpaceSelfTangent(x: 0, y: 5)
+  expectEqual(expectedGrad, actualGrad)
+}
+
+struct ProductSpaceOtherTangentTangentSpace : AdditiveArithmetic {
+  let x, y: Tracked<Float>
+}
+
+extension ProductSpaceOtherTangentTangentSpace : Differentiable {
+  typealias TangentVector = ProductSpaceOtherTangentTangentSpace
+}
+
+struct ProductSpaceOtherTangent {
+  var x, y: Tracked<Float>
+}
+
+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<Float> 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<Float>
+    @differentiable
+    var computed: Tracked<Float> {
+      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()

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<Float>) -> Tracked<Float>
+}
+func b(g: A) -> Tracked<Float> {
+  return gradient(at: 3) { x in g.a(x) }
+}
+
+struct B : A {
+  @differentiable(wrt: x)
+  func a(_ x: Tracked<Float>) -> Tracked<Float> { return x * 5 }
+}
+
+ExistentialTests.testWithLeakChecking("Existential method VJP") {
+  expectEqual(5.0, b(g: B()))
+}
+
+runAllTests()