[AD] Rewrite some of the tests with Tracked<Float> (2) (#27767)

* Leak checking in superset_adjoint.swift

* Leak checking in method.swift.

* Add a valueWithGradient function.

* Leak checking class_method.swift.

* Leak checking for existential.swift.

* Simplify  versions of gradient.

* Leak checking for e2e_differentiable_property.swift

* Leak checking for custom_derivatives.swift

* Some leak checking in array.swift

* Some leak checking in forward_mode_runtime.swift

* Formatting.

* Formatting update to test/AutoDiff/custom_derivatives.swift

Co-Authored-By: Richard Wei <[email protected]>

Author: bgogul

stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift

@@ -257,18 +257,27 @@ 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<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 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 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 {

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)