[AutoDiff upstream] Add Differentiable.withDerivative(_:). (swiftlang#30945)

Add `Differentiable.withDerivative(_:)`, a "derivative surgery" API.

`Differentiable.withDerivative(_:)` is an identity function returning `self`.
It takes a closure and applies it to the derivative of the return value, in
contexts where the return value is differentiated with respect to.

Author: dan-zheng

stdlib/public/Differentiation/DifferentiationUtilities.swift

@@ -17,6 +17,10 @@
 
 import Swift
 
+//===----------------------------------------------------------------------===//
+// Differentiable function creation
+//===----------------------------------------------------------------------===//
+
 /// Create a differentiable function from a vector-Jacobian products function.
 @inlinable
 public func differentiableFunction<T : Differentiable, R : Differentiable>(
@@ -70,6 +74,10 @@ public func differentiableFunction<T, U, V, R>(
     /*vjp*/ vjp)
 }
 
+//===----------------------------------------------------------------------===//
+// Derivative customization
+//===----------------------------------------------------------------------===//
+
 /// Returns `x` like an identity function. When used in a context where `x` is
 /// being differentiated with respect to, this function will not produce any 
 /// derivative at `x`.
@@ -91,6 +99,31 @@ public func withoutDerivative<T, R>(at x: T, in body: (T) -> R) -> R {
   body(x)
 }
 
+public extension Differentiable {
+  /// Applies the given closure to the derivative of `self`.
+  ///
+  /// Returns `self` like an identity function. When the return value is used in
+  /// a context where it is differentiated with respect to, applies the given
+  /// closure to the derivative of the return value.
+  @inlinable
+  @differentiable(wrt: self)
+  func withDerivative(_ body: @escaping (inout TangentVector) -> Void) -> Self {
+    return self
+  }
+
+  @inlinable
+  @derivative(of: withDerivative)
+  internal func _vjpWithDerivative(
+    _ body: @escaping (inout TangentVector) -> Void
+  ) -> (value: Self, pullback: (TangentVector) -> TangentVector) {
+    return (self, { grad in
+      var grad = grad
+      body(&grad)
+      return grad
+    })
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Diagnostics
 //===----------------------------------------------------------------------===//

test/AutoDiff/stdlib/derivative_customization.swift

@@ -0,0 +1,53 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import DifferentiationUnittest
+import StdlibUnittest
+
+var DerivativeCustomizationTests = TestSuite("DerivativeCustomization")
+
+DerivativeCustomizationTests.testWithLeakChecking("withDerivative") {
+  do {
+    var counter = 0
+    func callback(_ x: inout Tracked<Float>) { counter += 1 }
+    _ = gradient(at: 4) { (x: Tracked<Float>) -> Tracked<Float> in
+      // Non-active value should not be differentiated, so `callback` should
+      // not be called.
+      _ = x.withDerivative(callback)
+      return x.withDerivative(callback) + x.withDerivative(callback)
+    }
+    expectEqual(2, counter)
+  }
+
+  expectEqual(
+    30,
+    gradient(at: 4) { (x: Tracked<Float>) in
+      x.withDerivative { $0 = 10 } + x.withDerivative { $0 = 20 }
+    })
+}
+
+DerivativeCustomizationTests.testWithLeakChecking("withoutDerivative") {
+  expectEqual(
+    0,
+    gradient(at: Tracked<Float>(4)) { x -> Tracked<Float> in
+      withoutDerivative(at: x) { x in
+        x * x * x
+      }
+    })
+
+  expectEqual(
+    0,
+    gradient(at: Tracked<Float>(4)) { x -> Tracked<Float> in
+      let y = withoutDerivative(at: x)
+      return y * y * y
+    })
+
+  expectEqual(
+    2,
+    gradient(at: Tracked<Float>(4)) { x -> Tracked<Float> in
+      let y = withoutDerivative(at: x)
+      return x + y * y * y + x
+    })
+}
+
+runAllTests()