From fccb33a2c332ff2c3a940567c10d218a5cda0afc Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Wed, 16 Oct 2019 23:50:40 -0700
Subject: [PATCH 01/12] Leak checking in superset_adjoint.swift

---
 test/AutoDiff/superset_adjoint.swift | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/test/AutoDiff/superset_adjoint.swift b/test/AutoDiff/superset_adjoint.swift
index 361836f1f38df..6c475d4d11c66 100644
--- a/test/AutoDiff/superset_adjoint.swift
+++ b/test/AutoDiff/superset_adjoint.swift
@@ -51,7 +51,7 @@ SupersetVJPTests.test("ApplySubset") {
   func foo<T: Differentiable>(_ x: T, _ y: T, apply: @differentiable (T, T) -> T) -> T {
     return apply(x, y)
   }
-  expectEqual(1, gradient(at: Float(0)) { x in foo(x, 0) { $0 + $1 } })
+  expectEqual(1, gradient(at: Tracked<Float>(0)) { x in foo(x, 0) { $0 + $1 } })
 }
 
 // FIXME: The expression `(+) as @differentiable (Float, @nondiff Float) -> Float)`
@@ -63,16 +63,18 @@ SupersetVJPTests.test("ApplySubset") {
 //   expectEqual(Float(1), grad)
 // }
 
-SupersetVJPTests.test("IndirectResults") {
+SupersetVJPTests.testWithLeakChecking("IndirectResults") {
   @differentiable(wrt: (x, y), vjp: dx_T)
-  func x_T<T : Differentiable>(_ x: Float, _ y: T) -> Float {
+  func x_T<T : Differentiable>(_ x: Tracked<Float>, _ y: T) -> Tracked<Float> {
     if x > 1000 { return x }
     return x
   }
-  func dx_T<T : Differentiable>(_ x: Float, _ y: T) -> (Float, (Float) -> (Float, T.TangentVector)) {
+  func dx_T<T : Differentiable>(
+    _ x: Tracked<Float>, _ y: T
+  ) -> (Tracked<Float>, (Tracked<Float>) -> (Tracked<Float>, T.TangentVector)) {
     return (x_T(x, y), { v in (x * v, .zero) })
   }
-  expectEqual(2, gradient(at: 2) { x in x_T(x, Float(3)) })
+  expectEqual(2, gradient(at: 2) { x in x_T(x, Tracked<Float>(3)) })
 }
 
 runAllTests()

From e2b95df480919f1710ba92b6d6950c0794952dab Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Thu, 17 Oct 2019 00:09:32 -0700
Subject: [PATCH 02/12] Leak checking in method.swift.

---
 test/AutoDiff/method.swift | 163 +++++++++++++++++++++----------------
 1 file changed, 92 insertions(+), 71 deletions(-)

diff --git a/test/AutoDiff/method.swift b/test/AutoDiff/method.swift
index 7184ff8d24420..1b13317682057 100644
--- a/test/AutoDiff/method.swift
+++ b/test/AutoDiff/method.swift
@@ -2,54 +2,55 @@
 // REQUIRES: executable_test
 
 import StdlibUnittest
+import DifferentiationUnittest
 
 var MethodTests = TestSuite("Method")
 
 // ==== Tests with generated adjoint ====
 
 struct Parameter : Equatable {
-  private let storedX: Float
+  private let storedX: Tracked<Float>
   @differentiable(wrt: (self), jvp: jvpX, vjp: vjpX)
-  var x: Float {
+  var x: Tracked<Float> {
       return storedX
   }
   
-  init(x: Float) {
+  init(x: Tracked<Float>) {
     storedX = x
   }
 
-  func vjpX() -> (Float, (Float) -> Parameter) {
+  func vjpX() -> (Tracked<Float>, (Tracked<Float>) -> Parameter) {
     return (x, { dx in Parameter(x: dx) } )
   }
 
-  func jvpX() -> (Float, (Parameter) -> Float) {
+  func jvpX() -> (Tracked<Float>, (Parameter) -> Tracked<Float>) {
     return (x, { $0.x })
   }
 }
 
 extension Parameter {
-  func squared() -> Float {
+  func squared() -> Tracked<Float> {
     return x * x
   }
 
-  static func squared(p: Parameter) -> Float {
+  static func squared(p: Parameter) -> Tracked<Float> {
     return p.x * p.x
   }
 
-  func multiplied(with other: Float) -> Float {
+  func multiplied(with other: Tracked<Float>) -> Tracked<Float> {
     return x * other
   }
 
-  static func * (_ a: Parameter, _ b: Parameter) -> Float {
+  static func * (_ a: Parameter, _ b: Parameter) -> Tracked<Float> {
     return a.x * b.x
   }
 }
 
 extension Parameter : Differentiable, AdditiveArithmetic {
   typealias TangentVector = Parameter
-  typealias Scalar = Float
+  typealias Scalar = Tracked<Float>
   typealias Shape = ()
-  init(repeating repeatedValue: Float, shape: ()) {
+  init(repeating repeatedValue: Tracked<Float>, shape: ()) {
     self.init(x: repeatedValue)
   }
   static func + (lhs: Parameter, rhs: Parameter) -> Parameter {
@@ -64,32 +65,36 @@ extension Parameter : Differentiable, AdditiveArithmetic {
   static var zero: Parameter { return Parameter(x: 0) }
 }
 
-MethodTests.test("instance method with generated adjoint, called from differentated func") {
-  func f(_ p: Parameter) -> Float {
+MethodTests.testWithLeakChecking(
+  "instance method with generated adjoint, called from differentated func"
+) {
+  func f(_ p: Parameter) -> Tracked<Float> {
     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.test("instance method with generated adjoint, differentiated directly") {
+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) -> Float { p.squared() }
+  func g(p: Parameter) -> Tracked<Float> { 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.test("instance method with generated adjoint, wrt only self") {
-  func f(_ p: Parameter) -> Float {
+MethodTests.testWithLeakChecking("instance method with generated adjoint, wrt only self") {
+  func f(_ p: Parameter) -> Tracked<Float> {
     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.test("instance method with generated adjoint, wrt only non-self") {
-  func f(_ other: Float) -> Float {
+MethodTests.testWithLeakChecking("instance method with generated adjoint, wrt only non-self") {
+  func f(_ other: Tracked<Float>) -> Tracked<Float> {
     return 100 * Parameter(x: 200).multiplied(with: other)
   }
   expectEqual(100 * 200, gradient(at: 1, in: f))
@@ -98,14 +103,18 @@ MethodTests.test("instance method with generated adjoint, wrt only non-self") {
 
 // FIXME: Add a binary differential operator.
 //
-// MethodTests.test("instance method with generated adjoint, wrt self and non-self") {
-//   let g = #gradient({ (p: Parameter, o: Float) in p.multiplied(with: o) })
+// MethodTests.testWithLeakChecking(
+//   "instance method with generated adjoint, wrt self and non-self"
+// ) {
+//   let g = #gradient({ (p: Parameter, o: Tracked<Float>) 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.test("static method with generated adjoint, called from differentiated func") {
-  func f(_ p: Parameter) -> Float {
+MethodTests.testWithLeakChecking(
+  "static method with generated adjoint, called from differentiated func"
+) {
+  func f(_ p: Parameter) -> Tracked<Float> {
     return 100 * Parameter.squared(p: p)
   }
   expectEqual(Parameter(x: 4 * 100), gradient(at: Parameter(x: 2), in: f))
@@ -113,30 +122,32 @@ MethodTests.test("static method with generated adjoint, called from differentiat
 }
 
 // TODO(SR-8699): Fix this test.
-// MethodTests.test("static method with generated adjoint, differentiated directly") {
+// 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.test("static method with generated adjoint, wrt only first param") {
-  func f(_ p: Parameter) -> Float {
+MethodTests.testWithLeakChecking("static method with generated adjoint, wrt only first param") {
+  func f(_ p: Parameter) -> Tracked<Float> {
     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.test("static method with generated adjoint, wrt only second param") {
-  func f(_ p: Parameter) -> Float {
+MethodTests.testWithLeakChecking("static method with generated adjoint, wrt only second param") {
+  func f(_ p: Parameter) -> Tracked<Float> {
     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.test("static method with generated adjoint, wrt all params") {
-  func g(a: Parameter, b: Parameter) -> Float { a * b }
+MethodTests.testWithLeakChecking("static method with generated adjoint, wrt all params") {
+  func g(a: Parameter, b: Parameter) -> Tracked<Float> { 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)),
@@ -162,26 +173,26 @@ struct DiffWrtSelf : Differentiable {
 }
 
 struct CustomParameter : Equatable {
-  let storedX: Float
+  let storedX: Tracked<Float>
   @differentiable(wrt: (self), vjp: vjpX)
-  var x: Float {
+  var x: Tracked<Float> {
       return storedX
   }
   
-  init(x: Float) {
+  init(x: Tracked<Float>) {
     storedX = x
   }
 
-  func vjpX() -> (Float, (Float) -> CustomParameter) {
+  func vjpX() -> (Tracked<Float>, (Tracked<Float>) -> CustomParameter) {
     return (x, { dx in CustomParameter(x: dx) })
   }
 }
 
 extension CustomParameter : Differentiable, AdditiveArithmetic {
   typealias TangentVector = CustomParameter
-  typealias Scalar = Float
+  typealias Scalar = Tracked<Float>
   typealias Shape = ()
-  init(repeating repeatedValue: Float, shape: ()) {
+  init(repeating repeatedValue: Tracked<Float>, shape: ()) {
     self.init(x: repeatedValue)
   }
   static func + (lhs: CustomParameter, rhs: CustomParameter) -> CustomParameter {
@@ -196,28 +207,30 @@ extension CustomParameter : Differentiable, AdditiveArithmetic {
   static var zero: CustomParameter { return CustomParameter(x: 0) }
 }
 
-extension Float {
-  func clamped(to limits: ClosedRange<Float>) -> Float {
+extension Tracked where T : FloatingPoint {
+  func clamped(to limits: ClosedRange<Tracked<T>>) -> Tracked<T> {
     return min(max(self, limits.lowerBound), limits.upperBound)
   }
 }
 
 extension CustomParameter {
   @differentiable(wrt: (self), vjp: dSquared)
-  func squared() -> Float {
+  func squared() -> Tracked<Float> {
     return x * x
   }
 
-  func dSquared() -> (Float, (Float) -> CustomParameter) {
+  func dSquared() -> (Tracked<Float>, (Tracked<Float>) -> CustomParameter) {
     return (squared(), { [x] v in CustomParameter(x: (2 * x).clamped(to: -10.0...10.0) * v) })
   }
 
   @differentiable(vjp: dSquared)
-  static func squared(p: CustomParameter) -> Float {
+  static func squared(p: CustomParameter) -> Tracked<Float> {
     return p.x * p.x
   }
 
-  static func dSquared(_ p: CustomParameter) -> (Float, (Float) -> CustomParameter) {
+  static func dSquared(
+    _ p: CustomParameter
+  ) -> (Tracked<Float>, (Tracked<Float>) -> CustomParameter) {
     return (p.x * p.x, { v in CustomParameter(x: (2 * p.x).clamped(to: -10.0...10.0) * v) })
   }
 
@@ -225,79 +238,87 @@ extension CustomParameter {
   // parameters on the same func, so we define a copy of this func per adjoint.
 
   @differentiable(wrt: (self, other), vjp: dMultiplied_wrtAll)
-  func multiplied(with other: Float) -> Float {
+  func multiplied(with other: Tracked<Float>) -> Tracked<Float> {
     return x * other
   }
 
   @differentiable(wrt: (other), vjp: dMultiplied_wrtOther)
-  func multiplied_constSelf(with other: Float) -> Float {
+  func multiplied_constSelf(with other: Tracked<Float>) -> Tracked<Float> {
     return x * other
   }
 
   @differentiable(wrt: (self), vjp: dMultiplied_wrtSelf)
-  func multiplied_constOther(with other: Float) -> Float {
+  func multiplied_constOther(with other: Tracked<Float>) -> Tracked<Float> {
     return x * other
   }
 
-  func dMultiplied_wrtAll(with other: Float) -> (Float, (Float) -> (CustomParameter, Float)) {
+  func dMultiplied_wrtAll(
+    with other: Tracked<Float>
+  ) -> (Tracked<Float>, (Tracked<Float>) -> (CustomParameter, Tracked<Float>)) {
     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) })
   }
 
-  func dMultiplied_wrtOther(with other: Float) -> (Float, (Float) -> Float) {
+  func dMultiplied_wrtOther(
+    with other: Tracked<Float>
+  ) -> (Tracked<Float>, (Tracked<Float>) -> Tracked<Float>) {
     let (r, pb) = dMultiplied_wrtAll(with: other)
     return (r, { v in pb(v).1 })
   }
 
-  func dMultiplied_wrtSelf(with other: Float) -> (Float, (Float) -> CustomParameter) {
+  func dMultiplied_wrtSelf(
+    with other: Tracked<Float>
+  ) -> (Tracked<Float>, (Tracked<Float>) -> CustomParameter) {
     let (r, pb) = dMultiplied_wrtAll(with: other)
     return (r, { v in pb(v).0 })
   }
 
   @differentiable(vjp: dMultiply_wrtAll)
   static func multiply(_ lhs: CustomParameter, _ rhs: CustomParameter)
-      -> Float {
+      -> Tracked<Float> {
     return lhs.x * rhs.x
   }
 
   @differentiable(wrt: (rhs), vjp: dMultiply_wrtRhs)
-  static func multiply_constLhs(_ lhs: CustomParameter, _ rhs: CustomParameter) -> Float {
+  static func multiply_constLhs(_ lhs: CustomParameter, _ rhs: CustomParameter) -> Tracked<Float> {
     return lhs.x * rhs.x
   }
 
   static func dMultiply_wrtAll(_ lhs: CustomParameter,_ rhs: CustomParameter)
-      -> (Float, (Float) -> (CustomParameter, CustomParameter)) {
+      -> (Tracked<Float>, (Tracked<Float>) -> (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)) })
   }
 
   static func dMultiply_wrtRhs(_ lhs: CustomParameter, _ rhs: CustomParameter)
-      -> (Float, (Float) -> CustomParameter) {
+      -> (Tracked<Float>, (Tracked<Float>) -> CustomParameter) {
     let (r, pb) = dMultiply_wrtAll(lhs, rhs)
     return (r, { v in pb(v).1 })
   }
 }
 
-MethodTests.test("instance method with custom adjoint, called from differentated func") {
-  func f(_ p: CustomParameter) -> Float {
+MethodTests.testWithLeakChecking(
+  "instance method with custom adjoint, called from differentated func"
+) {
+  func f(_ p: CustomParameter) -> Tracked<Float> {
     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.test("instance method with generated adjoint, differentated directly") {
+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) -> Float { p.squared() }
+  func g(p: CustomParameter) -> Tracked<Float> { 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.test("static method with custom adjoint, called from differentated func") {
-  func f(_ p: CustomParameter) -> Float {
+MethodTests.testWithLeakChecking("static method with custom adjoint, called from differentated func") {
+  func f(_ p: CustomParameter) -> Tracked<Float> {
     return 100 * CustomParameter.squared(p: p)
   }
   expectEqual(CustomParameter(x: 4 * 100), gradient(at: CustomParameter(x: 2), in: f))
@@ -305,52 +326,52 @@ MethodTests.test("static method with custom adjoint, called from differentated f
 }
 
 // TODO(SR-8699): Fix this test.
-// MethodTests.test("static method with custom adjoint, differentiated directly") {
+// 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.test("instance method with custom adjoint, wrt only self") {
-  func f(_ p: CustomParameter) -> Float {
+MethodTests.testWithLeakChecking("instance method with custom adjoint, wrt only self") {
+  func f(_ p: CustomParameter) -> Tracked<Float> {
     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.test("instance method with custom adjoint, wrt only non-self") {
-  func f(_ other: Float) -> Float {
+MethodTests.testWithLeakChecking("instance method with custom adjoint, wrt only non-self") {
+  func f(_ other: Tracked<Float>) -> Tracked<Float> {
     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.test("instance method with custom adjoint, wrt self and non-self") {
-  func g(p: CustomParameter, o: Float) -> Float { p.multiplied(with: o) }
+MethodTests.testWithLeakChecking("instance method with custom adjoint, wrt self and non-self") {
+  func g(p: CustomParameter, o: Tracked<Float>) -> Tracked<Float> { 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.test("static method with custom adjoint, wrt only lhs") {
-  func f(_ p: CustomParameter) -> Float {
+MethodTests.testWithLeakChecking("static method with custom adjoint, wrt only lhs") {
+  func f(_ p: CustomParameter) -> Tracked<Float> {
     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.test("static method with custom adjoint, wrt only rhs") {
-  func f(_ p: CustomParameter) -> Float {
+MethodTests.testWithLeakChecking("static method with custom adjoint, wrt only rhs") {
+  func f(_ p: CustomParameter) -> Tracked<Float> {
     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.test("static method with custom adjoint, wrt all") {
-  func f(_ a: CustomParameter, _ b: CustomParameter) -> Float {
+MethodTests.testWithLeakChecking("static method with custom adjoint, wrt all") {
+  func f(_ a: CustomParameter, _ b: CustomParameter) -> Tracked<Float> {
     return CustomParameter.multiply(a, b)
   }
   expectEqual((CustomParameter(x: 5), CustomParameter(x: 10)),

From a8c9ca79242c58936fdc863cee35f9872d78dd2c Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Thu, 17 Oct 2019 15:01:07 -0700
Subject: [PATCH 03/12] Add a valueWithGradient function.

---
 .../DifferentiationUnittest/DifferentiationUnittest.swift  | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift b/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift
index 252f4f218c4d4..efd81b790bf7b 100644
--- a/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift
+++ b/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift
@@ -271,6 +271,13 @@ public func gradient(
   return pullback(at: x, y, in: f)(1)
 }
 
+public func valueWithGradient<T, U: FloatingPoint>(
+  at x: T, in f: @differentiable (T) -> Tracked<U>
+) -> (value: Tracked<U>, gradient: T.TangentVector) {
+  let (y, pullback) = valueWithPullback(at: x, in: f)
+  return (y, pullback(Tracked<U>(1)))
+}
+
 public extension Differentiable {
   @inlinable
   func gradient(

From 016a8bc4e43c251d77a1dbdd9002657e152e6f62 Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Thu, 17 Oct 2019 15:34:34 -0700
Subject: [PATCH 04/12] Leak checking class_method.swift.

---
 test/AutoDiff/class_method.swift | 168 ++++++++++++++++++-------------
 1 file changed, 97 insertions(+), 71 deletions(-)

diff --git a/test/AutoDiff/class_method.swift b/test/AutoDiff/class_method.swift
index c4f67e4b7937b..3661ca25650ca 100644
--- a/test/AutoDiff/class_method.swift
+++ b/test/AutoDiff/class_method.swift
@@ -2,56 +2,59 @@
 // REQUIRES: executable_test
 
 import StdlibUnittest
+import DifferentiationUnittest
 
 var ClassMethodTests = TestSuite("ClassMethods")
 
 ClassMethodTests.test("Final") {
   final class Final : Differentiable {
-    func method(_ x: Float) -> Float {
+    func method(_ x: Tracked<Float>) -> Tracked<Float> {
       return x * x
     }
   }
 
   for i in -5...5 {
-    expectEqual(Float(i) * 2, gradient(at: Float(i)) { x in Final().method(x) })
+    expectEqual(Tracked<Float>(Float(i * 2)), gradient(at: Tracked<Float>(Float(i))) {
+        x in Final().method(x)
+      })
   }
 }
 
 ClassMethodTests.test("Simple") {
   class Super {
     @differentiable(wrt: x, jvp: jvpf, vjp: vjpf)
-    func f(_ x: Float) -> Float {
+    func f(_ x: Tracked<Float>) -> Tracked<Float> {
       return 2 * x
     }
-    final func jvpf(_ x: Float) -> (Float, (Float) -> Float) {
+    final func jvpf(_ x: Tracked<Float>) -> (Tracked<Float>, (Tracked<Float>) -> Tracked<Float>) {
       return (f(x), { v in 2 * v })
     }
-    final func vjpf(_ x: Float) -> (Float, (Float) -> Float) {
+    final func vjpf(_ x: Tracked<Float>) -> (Tracked<Float>, (Tracked<Float>) -> Tracked<Float>) {
       return (f(x), { v in 2 * v })
     }
   }
 
   class SubOverride : Super {
     @differentiable(wrt: x)
-    override func f(_ x: Float) -> Float {
+    override func f(_ x: Tracked<Float>) -> Tracked<Float> {
       return 3 * x
     }
   }
 
   class SubOverrideCustomDerivatives : Super {
     @differentiable(wrt: x, jvp: jvpf2, vjp: vjpf2)
-    override func f(_ x: Float) -> Float {
+    override func f(_ x: Tracked<Float>) -> Tracked<Float> {
       return 3 * x
     }
-    final func jvpf2(_ x: Float) -> (Float, (Float) -> Float) {
+    final func jvpf2(_ x: Tracked<Float>) -> (Tracked<Float>, (Tracked<Float>) -> Tracked<Float>) {
       return (f(x), { v in 3 * v })
     }
-    final func vjpf2(_ x: Float) -> (Float, (Float) -> Float) {
+    final func vjpf2(_ x: Tracked<Float>) -> (Tracked<Float>, (Tracked<Float>) -> Tracked<Float>) {
       return (f(x), { v in 3 * v })
     }
   }
 
-  func classValueWithGradient(_ c: Super) -> (Float, Float) {
+  func classValueWithGradient(_ c: Super) -> (Tracked<Float>, Tracked<Float>) {
     return valueWithGradient(at: 1) { c.f($0) }
   }
   expectEqual((2, 2), classValueWithGradient(Super()))
@@ -61,28 +64,32 @@ ClassMethodTests.test("Simple") {
 
 ClassMethodTests.test("SimpleWrtSelf") {
   class Super : Differentiable {
-    var base: Float
+    var base: Tracked<Float>
     // FIXME(TF-648): Dummy to make `Super.AllDifferentiableVariables` be nontrivial.
-    var _nontrivial: [Float] = []
+    var _nontrivial: [Tracked<Float>] = []
 
     // TODO(TF-654): Uncomment attribute when differentiation supports class initializers.
     // TODO(TF-645): Remove `vjpInit` when differentiation supports `ref_element_addr`.
     // @differentiable(vjp: vjpInit)
-    required init(base: Float) {
+    required init(base: Tracked<Float>) {
       self.base = base
     }
-    static func vjpInit(base: Float) -> (Super, (TangentVector) -> Float) {
+    static func vjpInit(base: Tracked<Float>) -> (Super, (TangentVector) -> Tracked<Float>) {
       return (Super(base: base), { x in x.base })
     }
 
     @differentiable(wrt: (self, x), jvp: jvpf, vjp: vjpf)
-    func f(_ x: Float) -> Float {
+    func f(_ x: Tracked<Float>) -> Tracked<Float> {
       return base * x
     }
-    final func jvpf(_ x: Float) -> (Float, (TangentVector, Float) -> Float) {
+    final func jvpf(
+      _ x: Tracked<Float>
+    ) -> (Tracked<Float>, (TangentVector, Tracked<Float>) -> Tracked<Float>) {
       return (f(x), { (dself, dx) in dself.base * dx })
     }
-    final func vjpf(_ x: Float) -> (Float, (Float) -> (TangentVector, Float)) {
+    final func vjpf(
+      _ x: Tracked<Float>
+    ) -> (Tracked<Float>, (Tracked<Float>) -> (TangentVector, Tracked<Float>)) {
       let base = self.base
       return (f(x), { v in
         (TangentVector(base: v * x, _nontrivial: []), base * v)
@@ -92,7 +99,7 @@ ClassMethodTests.test("SimpleWrtSelf") {
 
   class SubOverride : Super {
     @differentiable(wrt: (self, x))
-    override func f(_ x: Float) -> Float {
+    override func f(_ x: Tracked<Float>) -> Tracked<Float> {
       return 3 * x
     }
   }
@@ -100,13 +107,13 @@ ClassMethodTests.test("SimpleWrtSelf") {
   class SubOverrideCustomDerivatives : Super {
     @differentiable(wrt: (self, x))
     @differentiable(wrt: x, jvp: jvpf2, vjp: vjpf2)
-    override func f(_ x: Float) -> Float {
+    override func f(_ x: Tracked<Float>) -> Tracked<Float> {
       return 3 * x
     }
-    final func jvpf2(_ x: Float) -> (Float, (Float) -> Float) {
+    final func jvpf2(_ x: Tracked<Float>) -> (Tracked<Float>, (Tracked<Float>) -> Tracked<Float>) {
       return (f(x), { v in 3 * v })
     }
-    final func vjpf2(_ x: Float) -> (Float, (Float) -> Float) {
+    final func vjpf2(_ x: Tracked<Float>) -> (Tracked<Float>, (Tracked<Float>) -> Tracked<Float>) {
       return (f(x), { v in 3 * v })
     }
   }
@@ -121,7 +128,7 @@ ClassMethodTests.test("SimpleWrtSelf") {
 
   // `valueWithGradient` is not used because nested tuples cannot be compared
   // with `expectEqual`.
-  func classGradient(_ c: Super) -> (Super.TangentVector, Float) {
+  func classGradient(_ c: Super) -> (Super.TangentVector, Tracked<Float>) {
     return gradient(at: c, 10) { c, x in c.f(x) }
   }
   expectEqual((Super.TangentVector(base: 10, _nontrivial: []), 2),
@@ -135,53 +142,66 @@ ClassMethodTests.test("SimpleWrtSelf") {
 ClassMethodTests.test("Generics") {
   class Super<T : Differentiable & FloatingPoint> where T == T.TangentVector {
     @differentiable(wrt: x, jvp: jvpf, vjp: vjpf)
-    func f(_ x: T) -> T {
-      return T(2) * x
+    func f(_ x: Tracked<T>) -> Tracked<T> {
+      return Tracked<T>(2) * x
     }
-    final func jvpf(_ x: T) -> (T, (T.TangentVector) -> T.TangentVector) {
-      return (f(x), { v in T(2) * v })
+    final func jvpf(
+      _ x: Tracked<T>
+    ) -> (Tracked<T>, (Tracked<T>.TangentVector) -> Tracked<T>.TangentVector) {
+      return (f(x), { v in Tracked<T>(2) * v })
     }
-    final func vjpf(_ x: T) -> (T, (T.TangentVector) -> T.TangentVector) {
-      return (f(x), { v in T(2) * v })
+    final func vjpf(
+      _ x: Tracked<T>
+    ) -> (Tracked<T>, (Tracked<T>.TangentVector) -> Tracked<T>.TangentVector) {
+      return (f(x), { v in Tracked<T>(2) * v })
     }
   }
 
   class SubOverride<T : Differentiable & FloatingPoint> : Super<T> where T == T.TangentVector {
     @differentiable(wrt: x)
-    override func f(_ x: T) -> T {
+    override func f(_ x: Tracked<T>) -> Tracked<T> {
       return x
     }
   }
 
   class SubSpecializeOverride : Super<Float> {
     @differentiable(wrt: x)
-    override func f(_ x: Float) -> Float {
+    override func f(_ x: Tracked<Float>) -> Tracked<Float> {
       return 3 * x
     }
   }
 
-  class SubOverrideCustomDerivatives<T : Differentiable & FloatingPoint> : Super<T> where T == T.TangentVector {
+  class SubOverrideCustomDerivatives<T : Differentiable & FloatingPoint> : Super<T>
+  where T == T.TangentVector {
     @differentiable(wrt: x, jvp: jvpf2, vjp: vjpf2)
-    override func f(_ x: T) -> T {
-      return T(3) * x
+    override func f(_ x: Tracked<T>) -> Tracked<T> {
+      return Tracked<T>(3) * x
     }
-    final func jvpf2(_ x: T) -> (T, (T.TangentVector) -> T.TangentVector) {
-      return (f(x), { v in T(3) * v })
+    final func jvpf2(
+      _ x: Tracked<T>
+    ) -> (Tracked<T>, (Tracked<T>.TangentVector) -> Tracked<T>.TangentVector) {
+      return (f(x), { v in Tracked<T>(3) * v })
     }
-    final func vjpf2(_ x: T) -> (T, (T.TangentVector) -> T.TangentVector) {
-      return (f(x), { v in T(3) * v })
+    final func vjpf2(
+      _ x: Tracked<T>
+    ) -> (Tracked<T>, (Tracked<T>.TangentVector) -> Tracked<T>.TangentVector) {
+      return (f(x), { v in Tracked<T>(3) * v })
     }
   }
 
   class SubSpecializeOverrideCustomDerivatives : Super<Float80> {
     @differentiable(wrt: x, jvp: jvpf2, vjp: vjpf2)
-    override func f(_ x: Float80) -> Float80 {
+    override func f(_ x: Tracked<Float80>) -> Tracked<Float80> {
       return 3 * x
     }
-    final func jvpf2(_ x: Float80) -> (Float80, (Float80) -> Float80) {
+    final func jvpf2(
+      _ x: Tracked<Float80>
+    ) -> (Tracked<Float80>, (Tracked<Float80>) -> Tracked<Float80>) {
       return (f(x), { v in 3 * v })
     }
-    final func vjpf2(_ x: Float80) -> (Float80, (Float80) -> Float80) {
+    final func vjpf2(
+      _ x: Tracked<Float80>
+    ) -> (Tracked<Float80>, (Tracked<Float80>) -> Tracked<Float80>) {
       return (f(x), { v in 3 * v })
     }
   }
@@ -189,7 +209,9 @@ ClassMethodTests.test("Generics") {
   func classValueWithGradient<T : Differentiable & FloatingPoint>(
     _ c: Super<T>
   ) -> (T, T) where T == T.TangentVector {
-    return valueWithGradient(at: T(1)) { c.f($0) }
+    let (x,y) =  valueWithGradient(at: Tracked<T>(1), in: {
+        c.f($0) })
+    return (x.value, y.value)
   }
   expectEqual((2, 2), classValueWithGradient(Super<Float>()))
   expectEqual((1, 1), classValueWithGradient(SubOverride<Float>()))
@@ -200,24 +222,24 @@ ClassMethodTests.test("Generics") {
 
 ClassMethodTests.test("Methods") {
   class Super : Differentiable {
-    var base: Float
+    var base: Tracked<Float>
     // Dummy to make `Super.AllDifferentiableVariables` be nontrivial.
-    var _nontrivial: [Float] = []
+    var _nontrivial: [Tracked<Float>] = []
 
     // TODO(TF-654): Uncomment attribute when differentiation supports class initializers.
     // TODO(TF-645): Remove `vjpInit` when differentiation supports `ref_element_addr`.
     // @differentiable(vjp: vjpInit)
-    init(base: Float) {
+    init(base: Tracked<Float>) {
       self.base = base
     }
-    static func vjpInit(base: Float) -> (Super, (TangentVector) -> Float) {
+    static func vjpInit(base: Tracked<Float>) -> (Super, (TangentVector) -> Tracked<Float>) {
       return (Super(base: base), { x in x.base })
     }
 
     @differentiable(vjp: vjpSquared)
-    func squared() -> Float { base * base }
+    func squared() -> Tracked<Float> { base * base }
 
-    final func vjpSquared() -> (Float, (Float) -> TangentVector) {
+    final func vjpSquared() -> (Tracked<Float>, (Tracked<Float>) -> TangentVector) {
       let base = self.base
       return (base * base, { v in
         TangentVector(base: 2 * base * v, _nontrivial: [])
@@ -227,8 +249,8 @@ ClassMethodTests.test("Methods") {
 
   class Sub1 : Super {
     @differentiable(vjp: vjpSquared2)
-    override func squared() -> Float { base * base }
-    final func vjpSquared2() -> (Float, (Float) -> TangentVector) {
+    override func squared() -> Tracked<Float> { base * base }
+    final func vjpSquared2() -> (Tracked<Float>, (Tracked<Float>) -> TangentVector) {
       let base = self.base
       return (base * base, { v in
         TangentVector(base: 2 * base * v, _nontrivial: [])
@@ -236,7 +258,7 @@ ClassMethodTests.test("Methods") {
     }
   }
 
-  func classValueWithGradient(_ c: Super) -> (Float, Super.TangentVector) {
+  func classValueWithGradient(_ c: Super) -> (Tracked<Float>, Super.TangentVector) {
     return valueWithGradient(at: c) { c in c.squared() }
   }
 
@@ -246,8 +268,8 @@ ClassMethodTests.test("Methods") {
   // TODO(TF-647): Handle `unchecked_ref_cast` in `Sub1.init` during pullback generation.
   // FIXME: `Super.init` VJP type mismatch for empty `Super.AllDifferentiableVariables`:
   // SIL verification failed: VJP type does not match expected VJP type
-  //   $@convention(method) (Float, @thick Super.Type) -> (@owned Super, @owned @callee_guaranteed (@guaranteed Super.AllDifferentiableVariables) -> Float)
-  //   $@convention(method) (Float, @owned Super) -> (@owned Super, @owned @callee_guaranteed (@guaranteed Super.AllDifferentiableVariables) -> Float)
+  //   $@convention(method) (Tracked<Float>, @thick Super.Type) -> (@owned Super, @owned @callee_guaranteed (@guaranteed Super.AllDifferentiableVariables) -> Tracked<Float>)
+  //   $@convention(method) (Tracked<Float>, @owned Super) -> (@owned Super, @owned @callee_guaranteed (@guaranteed Super.AllDifferentiableVariables) -> Tracked<Float>)
   // expectEqual(4, gradient(at: 2) { x in Sub1(base: x).squared() })
 
   expectEqual(Super.TangentVector(base: 4, _nontrivial: []),
@@ -258,20 +280,20 @@ ClassMethodTests.test("Methods") {
 
 ClassMethodTests.test("Properties") {
   class Super : Differentiable {
-    var base: Float
+    var base: Tracked<Float>
 
     // TODO(TF-654): Uncomment attribute when differentiation supports class initializers.
     // TODO(TF-645): Remove `vjpInit` when differentiation supports `ref_element_addr`.
     // @differentiable(vjp: vjpInit)
-    init(base: Float) { self.base = base }
-    static func vjpInit(base: Float) -> (Super, (TangentVector) -> Float) {
+    init(base: Tracked<Float>) { self.base = base }
+    static func vjpInit(base: Tracked<Float>) -> (Super, (TangentVector) -> Tracked<Float>) {
       return (Super(base: base), { x in x.base })
     }
 
     @differentiable(vjp: vjpSquared)
-    var squared: Float { base * base }
+    var squared: Tracked<Float> { base * base }
 
-    final func vjpSquared() -> (Float, (Float) -> TangentVector) {
+    final func vjpSquared() -> (Tracked<Float>, (Tracked<Float>) -> TangentVector) {
       let base = self.base
       return (base * base, { v in TangentVector(base: 2 * base * v) })
     }
@@ -281,13 +303,13 @@ ClassMethodTests.test("Properties") {
     // FIXME(TF-625): Crash due to `Super.AllDifferentiableVariables` abstraction pattern mismatch.
     // SIL verification failed: vtable entry for #<anonymous function>Super.squared!getter.1.jvp.1.S must be ABI-compatible
     //   ABI incompatible return values
-    //   @convention(method) (@guaranteed Super) -> (Float, @owned @callee_guaranteed (@guaranteed Super.AllDifferentiableVariables) -> Float)
-    //   @convention(method) (@guaranteed Sub1) -> (Float, @owned @callee_guaranteed (Super.AllDifferentiableVariables) -> Float)
+    //   @convention(method) (@guaranteed Super) -> (Tracked<Float>, @owned @callee_guaranteed (@guaranteed Super.AllDifferentiableVariables) -> Tracked<Float>)
+    //   @convention(method) (@guaranteed Sub1) -> (Tracked<Float>, @owned @callee_guaranteed (Super.AllDifferentiableVariables) -> Tracked<Float>)
     // @differentiable
-    // override var squared: Float { base * base }
+    // override var squared: Tracked<Float> { base * base }
   }
 
-  func classValueWithGradient(_ c: Super) -> (Float, Super.TangentVector) {
+  func classValueWithGradient(_ c: Super) -> (Tracked<Float>, Super.TangentVector) {
     return valueWithGradient(at: c) { c in c.squared }
   }
 
@@ -299,38 +321,42 @@ ClassMethodTests.test("Properties") {
 
 ClassMethodTests.test("Capturing") {
   class Multiplier {
-    var coefficient: Float
-    init(_ coefficient: Float) {
+    var coefficient: Tracked<Float>
+    init(_ coefficient: Tracked<Float>) {
       self.coefficient = coefficient
     }
 
     // Case 1: generated VJP.
     @differentiable
-    func apply(to x: Float) -> Float {
+    func apply(to x: Tracked<Float>) -> Tracked<Float> {
       return coefficient * x
     }
 
     // Case 2: custom VJP capturing `self`.
     @differentiable(wrt: (x), vjp: vjpApply2)
-    func apply2(to x: Float) -> Float {
+    func apply2(to x: Tracked<Float>) -> Tracked<Float> {
       return coefficient * x
     }
-    final func vjpApply2(to x: Float) -> (Float, (Float) -> Float) {
+    final func vjpApply2(
+      to x: Tracked<Float>
+    ) -> (Tracked<Float>, (Tracked<Float>) -> Tracked<Float>) {
       return (coefficient * x, { v in self.coefficient * v })
     }
 
     // Case 3: custom VJP capturing `self.coefficient`.
     @differentiable(wrt: x, vjp: vjpApply3)
-    func apply3(to x: Float) -> Float {
+    func apply3(to x: Tracked<Float>) -> Tracked<Float> {
       return coefficient * x
     }
-    final func vjpApply3(to x: Float) -> (Float, (Float) -> Float) {
+    final func vjpApply3(
+      to x: Tracked<Float>
+    ) -> (Tracked<Float>, (Tracked<Float>) -> Tracked<Float>) {
       let coefficient = self.coefficient
       return (coefficient * x, { v in coefficient * v })
     }
   }
 
-  func f(_ x: Float) -> Float {
+  func f(_ x: Tracked<Float>) -> Tracked<Float> {
     let m = Multiplier(10)
     let result = m.apply(to: x)
     m.coefficient += 1
@@ -338,7 +364,7 @@ ClassMethodTests.test("Capturing") {
   }
   expectEqual(10, gradient(at: 1, in: f))
 
-  func f2(_ x: Float) -> Float {
+  func f2(_ x: Tracked<Float>) -> Tracked<Float> {
     let m = Multiplier(10)
     let result = m.apply2(to: x)
     m.coefficient += 1
@@ -346,7 +372,7 @@ ClassMethodTests.test("Capturing") {
   }
   expectEqual(11, gradient(at: 1, in: f2))
 
-  func f3(_ x: Float) -> Float {
+  func f3(_ x: Tracked<Float>) -> Tracked<Float> {
     let m = Multiplier(10)
     let result = m.apply3(to: x)
     m.coefficient += 1

From f8291cb792e1b69cc8d12a29d29b3ce4cc434ca6 Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Thu, 17 Oct 2019 15:53:45 -0700
Subject: [PATCH 05/12] Leak checking for existential.swift.

---
 test/AutoDiff/existential.swift | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/test/AutoDiff/existential.swift b/test/AutoDiff/existential.swift
index 23665376d3f77..ba69eea19f93f 100644
--- a/test/AutoDiff/existential.swift
+++ b/test/AutoDiff/existential.swift
@@ -2,21 +2,22 @@
 // REQUIRES: executable_test
 
 import StdlibUnittest
+import DifferentiationUnittest
 
 var ExistentialTests = TestSuite("Existential")
 
 protocol A {
   @differentiable(wrt: x)
-  func a(_ x: Float) -> Float
+  func a(_ x: Tracked<Float>) -> Tracked<Float>
 }
-func b(g: A) -> Float { return (3.0 as Float).gradient() { x in g.a(x) } }
+func b(g: A) -> Tracked<Float> { return (3.0 as Tracked<Float>).gradient() { x in g.a(x) } }
 
 struct B : A {
   @differentiable(wrt: x)
-  func a(_ x: Float) -> Float { return x * 5.0 }
+  func a(_ x: Tracked<Float>) -> Tracked<Float> { return x * 5.0 }
 }
 
-ExistentialTests.test("vjp/adjoint constructed with existentials.") {
+ExistentialTests.testWithLeakChecking("vjp/adjoint constructed with existentials.") {
   expectEqual(5.0, b(g: B()))
 }
 

From 237afc568c444ca4d9ad4ca5ad9907dc0c364db4 Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Thu, 17 Oct 2019 16:10:05 -0700
Subject: [PATCH 06/12] Simplify  versions of gradient.

---
 .../DifferentiationUnittest.swift             | 24 ++++++++++---------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift b/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift
index efd81b790bf7b..9c223d890a2bf 100644
--- a/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift
+++ b/stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift
@@ -257,21 +257,23 @@ extension Tracked where T : Differentiable & FloatingPoint, T == T.TangentVector
   }
 }
 
-// Differential operators for `Tracked<Float>`.
-public func gradient(
-  at x: Tracked<Float>, in f: @differentiable (Tracked<Float>) -> Tracked<Float>
-) -> Tracked<Float> {
-  return pullback(at: x, in: f)(1)
+// Differential operators for `Tracked<T>`.
+
+public func gradient<T, U>(
+  at x: T, in f: @differentiable (T) -> Tracked<U>
+) -> T.TangentVector
+where U : FloatingPoint, U.TangentVector == U {
+  return pullback(at: x, in: f)(Tracked<U>(1))
 }
 
-public func gradient(
-  at x: Tracked<Float>, _ y: Tracked<Float>,
-  in f: @differentiable (Tracked<Float>, Tracked<Float>) -> Tracked<Float>
-) -> (Tracked<Float>, Tracked<Float>) {
-  return pullback(at: x, y, in: f)(1)
+public func gradient<T, U, R>(
+  at x: T, _ y: U, in f: @differentiable (T, U) -> Tracked<R>
+) -> (T.TangentVector, U.TangentVector)
+  where R : FloatingPoint, R.TangentVector == R {
+  return pullback(at: x, y, in: f)(Tracked<R>(1))
 }
 
-public func valueWithGradient<T, U: FloatingPoint>(
+public func valueWithGradient<T, U : FloatingPoint>(
   at x: T, in f: @differentiable (T) -> Tracked<U>
 ) -> (value: Tracked<U>, gradient: T.TangentVector) {
   let (y, pullback) = valueWithPullback(at: x, in: f)

From 03704df02a98ec8593cf24b31fe8e90987f4776d Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Thu, 17 Oct 2019 16:10:22 -0700
Subject: [PATCH 07/12] Leak checking for e2e_differentiable_property.swift

---
 .../e2e_differentiable_property.swift         | 49 ++++++++++---------
 1 file changed, 26 insertions(+), 23 deletions(-)

diff --git a/test/AutoDiff/e2e_differentiable_property.swift b/test/AutoDiff/e2e_differentiable_property.swift
index 521776282aafe..489c94ee9ed5e 100644
--- a/test/AutoDiff/e2e_differentiable_property.swift
+++ b/test/AutoDiff/e2e_differentiable_property.swift
@@ -5,11 +5,12 @@
 // VJPs for the properties specified in various ways.
 
 import StdlibUnittest
+import DifferentiationUnittest
 
 var E2EDifferentiablePropertyTests = TestSuite("E2EDifferentiableProperty")
 
 struct TangentSpace : AdditiveArithmetic {
-  let x, y: Float
+  let x, y: Tracked<Float>
 }
 
 extension TangentSpace : Differentiable {
@@ -18,22 +19,22 @@ extension TangentSpace : Differentiable {
 
 struct Space {
   /// `x` is a computed property with a custom vjp.
-  var x: Float {
+  var x: Tracked<Float> {
     @differentiable(vjp: vjpX)
     get { storedX }
     set { storedX = newValue }
   }
 
-  func vjpX() -> (Float, (Float) -> TangentSpace) {
+  func vjpX() -> (Tracked<Float>, (Tracked<Float>) -> TangentSpace) {
     return (x, { v in TangentSpace(x: v, y: 0) } )
   }
 
-  private var storedX: Float
+  private var storedX: Tracked<Float>
 
   @differentiable
-  var y: Float
+  var y: Tracked<Float>
 
-  init(x: Float, y: Float) {
+  init(x: Tracked<Float>, y: Tracked<Float>) {
     self.storedX = x
     self.y = y
   }
@@ -47,16 +48,16 @@ extension Space : Differentiable {
   }
 }
 
-E2EDifferentiablePropertyTests.test("computed property") {
-  let actualGrad = gradient(at: Space(x: 0, y: 0)) { (point: Space) -> Float in
+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.test("stored property") {
-  let actualGrad = gradient(at: Space(x: 0, y: 0)) { (point: Space) -> Float in
+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)
@@ -73,24 +74,24 @@ struct GenericMemberWrapper<T : Differentiable> : Differentiable {
   }
 }
 
-E2EDifferentiablePropertyTests.test("generic stored property") {
-  let actualGrad = gradient(at: GenericMemberWrapper<Float>(x: 1)) { point in
+E2EDifferentiablePropertyTests.testWithLeakChecking("generic stored property") {
+  let actualGrad = gradient(at: GenericMemberWrapper<Tracked<Float>>(x: 1)) { point in
     return 2 * point.x
   }
-  let expectedGrad = GenericMemberWrapper<Float>.TangentVector(x: 2)
+  let expectedGrad = GenericMemberWrapper<Tracked<Float>>.TangentVector(x: 2)
   expectEqual(expectedGrad, actualGrad)
 }
 
 struct ProductSpaceSelfTangent : AdditiveArithmetic {
-  let x, y: Float
+  let x, y: Tracked<Float>
 }
 
 extension ProductSpaceSelfTangent : Differentiable {
   typealias TangentVector = ProductSpaceSelfTangent
 }
 
-E2EDifferentiablePropertyTests.test("fieldwise product space, self tangent") {
-  let actualGrad = gradient(at: ProductSpaceSelfTangent(x: 0, y: 0)) { (point: ProductSpaceSelfTangent) -> Float in
+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)
@@ -98,7 +99,7 @@ E2EDifferentiablePropertyTests.test("fieldwise product space, self tangent") {
 }
 
 struct ProductSpaceOtherTangentTangentSpace : AdditiveArithmetic {
-  let x, y: Float
+  let x, y: Tracked<Float>
 }
 
 extension ProductSpaceOtherTangentTangentSpace : Differentiable {
@@ -106,7 +107,7 @@ extension ProductSpaceOtherTangentTangentSpace : Differentiable {
 }
 
 struct ProductSpaceOtherTangent {
-  var x, y: Float
+  var x, y: Tracked<Float>
 }
 
 extension ProductSpaceOtherTangent : Differentiable {
@@ -117,19 +118,21 @@ extension ProductSpaceOtherTangent : Differentiable {
   }
 }
 
-E2EDifferentiablePropertyTests.test("fieldwise product space, other tangent") {
-  let actualGrad = gradient(at: ProductSpaceOtherTangent(x: 0, y: 0)) { (point: ProductSpaceOtherTangent) -> Float in
+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.test("computed property") {
+E2EDifferentiablePropertyTests.testWithLeakChecking("computed property") {
   struct TF_544 : Differentiable {
-    var value: Float
+    var value: Tracked<Float>
     @differentiable
-    var computed: Float {
+    var computed: Tracked<Float> {
       get { value }
       set { value = newValue }
     }

From 41a9f93a0c83528702a040038f10e23d70580f77 Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Thu, 17 Oct 2019 16:18:35 -0700
Subject: [PATCH 08/12] Leak checking for custom_derivatives.swift

---
 test/AutoDiff/custom_derivatives.swift | 35 +++++++++++++-------------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/test/AutoDiff/custom_derivatives.swift b/test/AutoDiff/custom_derivatives.swift
index 7a894f37b0f1b..d1f2669bf0e18 100644
--- a/test/AutoDiff/custom_derivatives.swift
+++ b/test/AutoDiff/custom_derivatives.swift
@@ -7,25 +7,26 @@ import Darwin.C
 #else
 import Glibc
 #endif
+import DifferentiationUnittest
 
 var CustomDerivativesTests = TestSuite("CustomDerivatives")
 
 // Specify non-differentiable functions.
 // These will be wrapped in `differentiableFunction` and tested.
 
-func unary(_ x: Float) -> Float {
+func unary(_ x: Tracked<Float>) -> Tracked<Float> {
   var x = x
   x *= 2
   return x
 }
 
-func binary(_ x: Float, _ y: Float) -> Float {
+func binary(_ x: Tracked<Float>, _ y: Tracked<Float>) -> Tracked<Float> {
   var x = x
   x *= y
   return x
 }
 
-CustomDerivativesTests.test("differentiableFunction-unary") {
+CustomDerivativesTests.testWithLeakChecking("differentiableFunction-unary") {
   let diffableUnary = differentiableFunction { x in
     (value: unary(x), pullback: { v in v * x * 2 })
   }
@@ -35,7 +36,7 @@ CustomDerivativesTests.test("differentiableFunction-unary") {
   expectEqual(40, gradient(at: 10, in: { diffableUnary($0) * 2 }))
 }
 
-CustomDerivativesTests.test("differentiableFunction-binary") {
+CustomDerivativesTests.testWithLeakChecking("differentiableFunction-binary") {
   let diffableBinary = differentiableFunction { (x, y) in
     (value: binary(x, y), pullback: { v in (v * y, v * x) })
   }
@@ -45,14 +46,14 @@ CustomDerivativesTests.test("differentiableFunction-binary") {
   expectEqual((20, 10), gradient(at: 5, 10, in: { diffableBinary($0, $1) * 2 }))
 }
 
-CustomDerivativesTests.test("Checkpointing") {
+CustomDerivativesTests.testWithLeakChecking("Checkpointing") {
   var count = 0
-  func cube(_ x: Float) -> Float {
+  func cube(_ x: Tracked<Float>) -> Tracked<Float> {
     count += 1
     return x * x * x
   }
   // Test the top-level function variant of the checkpointing API.
-  expectEqual(324, gradient(at: 3) { (x: Float) -> Float in
+  expectEqual(324, gradient(at: 3) { (x: Tracked<Float>) -> Tracked<Float> in
     expectEqual(0, count)
     let y = withRecomputationInPullbacks(cube)(x)
     expectEqual(1, count)
@@ -61,7 +62,7 @@ CustomDerivativesTests.test("Checkpointing") {
   expectEqual(2, count)
   // Reset and test the method variant.
   count = 0
-  expectEqual(324, gradient(at: 3) { (x: Float) -> Float in
+  expectEqual(324, gradient(at: 3) { (x: Tracked<Float>) -> Tracked<Float> in
     expectEqual(0, count)
     let y = x.withRecomputationInPullbacks(cube)
     expectEqual(1, count)
@@ -70,10 +71,10 @@ CustomDerivativesTests.test("Checkpointing") {
   expectEqual(2, count)
 }
 
-CustomDerivativesTests.test("SumOfGradPieces") {
-  var grad: Float = 0
-  func addToGrad(_ x: inout Float) { grad += x }
-  _ = gradient(at: 4) { (x: Float) in
+CustomDerivativesTests.testWithLeakChecking("SumOfGradPieces") {
+  var grad: Tracked<Float> = 0
+  func addToGrad(_ x: inout Tracked<Float>) { grad += x }
+  _ = gradient(at: 4) { (x: Tracked<Float>) in
     x.withDerivative(addToGrad)
       * x.withDerivative(addToGrad)
         * x.withDerivative(addToGrad)
@@ -81,16 +82,16 @@ CustomDerivativesTests.test("SumOfGradPieces") {
   expectEqual(48, grad)
 }
 
-CustomDerivativesTests.test("ModifyGradientOfSum") {
-  expectEqual(30, gradient(at: 4) { (x: Float) in
+CustomDerivativesTests.testWithLeakChecking("ModifyGradientOfSum") {
+  expectEqual(30, gradient(at: 4) { (x: Tracked<Float>) in
     x.withDerivative { $0 *= 10 } + x.withDerivative { $0 *= 20 }
   })
 }
 
-CustomDerivativesTests.test("WithoutDerivative") {
-  expectEqual(0, gradient(at: Float(4)) { x in
+CustomDerivativesTests.testWithLeakChecking("WithoutDerivative") {
+  expectEqual(0, gradient(at: Tracked<Float>(4)) { x in
     withoutDerivative(at: x) { x in
-      sinf(x) + cosf(x)
+        Tracked<Float>(sinf(x.value) + cosf(x.value))
     }
   })
 }

From 0f2711a66e1a3376f1c5cf92069170f6ef0ed174 Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Thu, 17 Oct 2019 16:32:38 -0700
Subject: [PATCH 09/12] Some leak checking in array.swift

---
 test/AutoDiff/array.swift | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/AutoDiff/array.swift b/test/AutoDiff/array.swift
index 2a79888a8c673..0da15636e28a7 100644
--- a/test/AutoDiff/array.swift
+++ b/test/AutoDiff/array.swift
@@ -28,7 +28,7 @@ ArrayAutoDiffTests.test("ArraySubscript") {
     gradient(at: [2, 3, 4, 5, 6, 7], in: sumFirstThree))
 }
 
-ArrayAutoDiffTests.test("ArrayLiteral") {
+ArrayAutoDiffTests.testWithLeakChecking("ArrayLiteral") {
   func twoElementLiteral(_ x: Tracked<Float>, _ y: Tracked<Float>) -> [Tracked<Float>] {
     return [x, y]
   }
@@ -90,7 +90,7 @@ ArrayAutoDiffTests.test("ArrayConcat") {
       in: sumFirstThreeConcatted))
 }
 
-ArrayAutoDiffTests.test("Array.init(repeating:count:)") {
+ArrayAutoDiffTests.testWithLeakChecking("Array.init(repeating:count:)") {
   @differentiable
   func repeating(_ x: Tracked<Float>) -> [Tracked<Float>] {
     Array(repeating: x, count: 10)

From 4e052540f68591cbfc26f63bbe7298c6996b4fd0 Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Thu, 17 Oct 2019 16:30:51 -0700
Subject: [PATCH 10/12] Some leak checking in forward_mode_runtime.swift

---
 test/AutoDiff/forward_mode_runtime.swift | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/test/AutoDiff/forward_mode_runtime.swift b/test/AutoDiff/forward_mode_runtime.swift
index 89c1ef6e06f62..49e1afb996043 100644
--- a/test/AutoDiff/forward_mode_runtime.swift
+++ b/test/AutoDiff/forward_mode_runtime.swift
@@ -205,7 +205,7 @@ ForwardModeTests.test("complexBinaryMethod") {
 // Tracked struct
 //===----------------------------------------------------------------------===//
 
-ForwardModeTests.test("TrackedIdentity") {
+ForwardModeTests.testWithLeakChecking("TrackedIdentity") {
   func identity(x: Tracked<Float>) -> Tracked<Float> {
     return x
   }
@@ -214,7 +214,7 @@ ForwardModeTests.test("TrackedIdentity") {
   expectEqual(1, differential(1))
 }
 
-ForwardModeTests.test("TrackedAddition") {
+ForwardModeTests.testWithLeakChecking("TrackedAddition") {
   func add(x: Tracked<Float>, y: Tracked<Float>) -> Tracked<Float> {
     return x + y
   }
@@ -223,7 +223,7 @@ ForwardModeTests.test("TrackedAddition") {
   expectEqual(2, differential(1, 1))
 }
 
-ForwardModeTests.test("TrackedDivision") {
+ForwardModeTests.testWithLeakChecking("TrackedDivision") {
   func divide(x: Tracked<Float>, y: Tracked<Float>) -> Tracked<Float> {
     return x / y
   }
@@ -232,7 +232,7 @@ ForwardModeTests.test("TrackedDivision") {
   expectEqual(-0.2, differential(1, 1))
 }
 
-ForwardModeTests.test("TrackedMultipleMultiplication") {
+ForwardModeTests.testWithLeakChecking("TrackedMultipleMultiplication") {
   func add(x: Tracked<Float>, y: Tracked<Float>) -> Tracked<Float> {
     return x * y * x
   }
@@ -242,7 +242,7 @@ ForwardModeTests.test("TrackedMultipleMultiplication") {
   expectEqual(56, differential(1, 1))
 }
 
-ForwardModeTests.test("TrackedWithLets") {
+ForwardModeTests.testWithLeakChecking("TrackedWithLets") {
   func add(x: Tracked<Float>, y: Tracked<Float>) -> Tracked<Float> {
     let a = x + y
     let b = a * a // (x+y)^2
@@ -627,7 +627,7 @@ ForwardModeTests.test("GenericTrackedBinaryVars") {
   expectEqual(36, differential(1, 1))
 }
 
-ForwardModeTests.test("TrackedDifferentiableFuncType") {
+ForwardModeTests.testWithLeakChecking("TrackedDifferentiableFuncType") {
   func valAndDeriv(
     f: @escaping @differentiable (Tracked<Float>) -> Tracked<Float>
   ) -> (Tracked<Float>, Tracked<Float>) {

From 6b170e0a2ab2fe8d72169af0ae40f8bfe8e71d0d Mon Sep 17 00:00:00 2001
From: Dan Zheng <danielzheng@google.com>
Date: Thu, 17 Oct 2019 16:44:01 -0700
Subject: [PATCH 11/12] Formatting.

---
 test/AutoDiff/class_method.swift | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/test/AutoDiff/class_method.swift b/test/AutoDiff/class_method.swift
index 3661ca25650ca..78484a0fc3f83 100644
--- a/test/AutoDiff/class_method.swift
+++ b/test/AutoDiff/class_method.swift
@@ -14,9 +14,10 @@ ClassMethodTests.test("Final") {
   }
 
   for i in -5...5 {
-    expectEqual(Tracked<Float>(Float(i * 2)), gradient(at: Tracked<Float>(Float(i))) {
-        x in Final().method(x)
-      })
+  expectEqual(Tracked<Float>(Float(i * 2)),
+              gradient(at: Tracked<Float>(Float(i))) {
+                x in Final().method(x)
+              })
   }
 }
 

From 4750624410913e5010322019cdbe20cd63d6179a Mon Sep 17 00:00:00 2001
From: Gogul Balakrishnan <bgogul@google.com>
Date: Fri, 18 Oct 2019 08:39:14 -0700
Subject: [PATCH 12/12] Formatting update to
 test/AutoDiff/custom_derivatives.swift

Co-Authored-By: Richard Wei <rxwei@google.com>
---
 test/AutoDiff/custom_derivatives.swift | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/AutoDiff/custom_derivatives.swift b/test/AutoDiff/custom_derivatives.swift
index d1f2669bf0e18..676548a437255 100644
--- a/test/AutoDiff/custom_derivatives.swift
+++ b/test/AutoDiff/custom_derivatives.swift
@@ -91,7 +91,7 @@ CustomDerivativesTests.testWithLeakChecking("ModifyGradientOfSum") {
 CustomDerivativesTests.testWithLeakChecking("WithoutDerivative") {
   expectEqual(0, gradient(at: Tracked<Float>(4)) { x in
     withoutDerivative(at: x) { x in
-        Tracked<Float>(sinf(x.value) + cosf(x.value))
+      Tracked<Float>(sinf(x.value) + cosf(x.value))
     }
   })
 }