[AutoDiff] Enable some checkpointing APIs that became unblocked. (#22668)

This completes the set of checkpointing APIs introduced in #22033.

Author: rxwei

stdlib/public/core/AutoDiff.swift

@@ -156,24 +156,22 @@ public func withRecomputationInPullbacks<T, U>(
   }
 }
 
-// FIXME: The method variant produces a zero cotangent. Need to investigate.
-//
-// public extension Differentiable {
-//   @inlinable
-//   @differentiable(wrt: self, vjp: _vjp_withRecomputationInPullbacks)
-//   func withRecomputationInPullbacks<Result : Differentiable>(
-//     _ body: @escaping @differentiable (Self) -> Result
-//   ) -> Result {
-//     return body(self)
-//   }
-// 
-//   @usableFromInline
-//   internal func _vjp_withRecomputationInPullbacks<Result : Differentiable>(
-//     _ body: @escaping @differentiable (Self) -> Result
-//   ) -> (Result, (Result.CotangentVector) -> CotangentVector) {
-//     return valueWithPullback(in: Swift.withRecomputationInPullbacks(body))
-//   }
-// }
+public extension Differentiable {
+  @inlinable
+  @differentiable(wrt: self, vjp: _vjp_withRecomputationInPullbacks)
+  func withRecomputationInPullbacks<Result : Differentiable>(
+    _ body: @escaping @differentiable (Self) -> Result
+  ) -> Result {
+    return body(self)
+  }
+
+  @usableFromInline
+  internal func _vjp_withRecomputationInPullbacks<Result : Differentiable>(
+    _ body: @escaping @differentiable (Self) -> Result
+  ) -> (Result, (Result.CotangentVector) -> CotangentVector) {
+    return valueWithPullback(in: Swift.withRecomputationInPullbacks(body))
+  }
+}
 
 //===----------------------------------------------------------------------===//
 // Method-style differential operators

test/AutoDiff/custom_derivatives.swift

@@ -60,15 +60,14 @@ CustomDerivativesTests.test("Checkpointing") {
   })
   expectEqual(2, count)
   // Reset and test the method variant.
-  // FIXME: The method variant produces a zero cotangent. Need to investigate.
-  // count = 0
-  // expectEqual(324, gradient(at: 3) { (x: Float) -> Float in
-  //   expectEqual(0, count)
-  //   let y = x.withRecomputationInPullbacks(f)
-  //   expectEqual(1, count)
-  //   return y * 3 * x
-  // })
-  // expectEqual(2, count)
+  count = 0
+  expectEqual(324, gradient(at: 3) { (x: Float) -> Float in
+    expectEqual(0, count)
+    let y = x.withRecomputationInPullbacks(f)
+    expectEqual(1, count)
+    return y * 3 * x
+  })
+  expectEqual(2, count)
 }
 
 runAllTests()