Some TF stdlib fixes (#22112)

* [TF API] Fix VJPs for sum and mean

The VJP registered for the squeezingAxes variant of each of those
methods is in fact the correct VJP for the alongAxes variant; the VJP
for the squeezingAxes variant requires inserting an additional
dimension, which is left for future implementation.

* [TF API] Add Tensor.variance

* [TF API] Mark some more reduction methods `@differentiable where`

* [TF API] Remove tests for broken sum/mean VJPs

* [TF] check out branch of swift-apis

* [AutoDiff] Fix SILGen for `AutoDiffFunctionExtractExpr` so that it won't leak owned values.

Author: jekbradbury

lib/SILGen/SILGenExpr.cpp

@@ -5390,7 +5390,7 @@ RValue RValueEmitter::visitAutoDiffFunctionExtractOriginalExpr(
   auto diffFunc = SGF.emitRValueAsSingleValue(E->getSubExpr());
   auto *orig = SGF.B.createAutoDiffFunctionExtractOriginal(
       E, diffFunc.forward(SGF));
-  return RValue(SGF, E, ManagedValue::forUnmanaged(orig));
+  return RValue(SGF, E, SGF.emitManagedRValueWithCleanup(orig));
 }
 
 RValue RValueEmitter::visitTapExpr(TapExpr *E, SGFContext C) {

stdlib/public/TensorFlow/Gradients.swift

@@ -559,16 +559,16 @@ extension Tensor where Scalar : Differentiable & FloatingPoint {
   }
 
   @inlinable
-  func _vjpMean(squeezingAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
-    let value = mean(squeezingAxes: axes)
+  func _vjpMean(alongAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
+    let value = mean(alongAxes: axes)
     return (value, { [shape = shapeTensor, count = scalarCountTensor] in
       $0.broadcast(toShape: shape) / Tensor(count)
     })
   }
 
   @inlinable
-  func _vjpSum(squeezingAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
-    let value = sum(squeezingAxes: axes)
+  func _vjpSum(alongAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
+    let value = sum(alongAxes: axes)
     return (value, { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
   }
 }

stdlib/public/TensorFlow/Ops.swift

@@ -1225,10 +1225,6 @@ public extension Tensor where Scalar : Numeric {
   /// - Parameter axes: The dimensions to reduce.
   /// - Precondition: Each value in `axes` must be in the range `-rank...rank`.
   @inlinable @inline(__always)
-  @differentiable(
-    wrt: self, vjp: _vjpMean(squeezingAxes:)
-    where Scalar : Differentiable & FloatingPoint
-  )
   func mean(squeezingAxes axes: [Int32]) -> Tensor {
     return Raw.mean(self, reductionIndices: Tensor<Int32>(axes),
                     keepDims: false)
@@ -1239,10 +1235,6 @@ public extension Tensor where Scalar : Numeric {
   /// - Parameter axes: The dimensions to reduce.
   /// - Precondition: Each value in `axes` must be in the range `-rank...rank`.
   @inlinable @inline(__always)
-  @differentiable(
-    wrt: self, vjp: _vjpMean(squeezingAxes:)
-    where Scalar : Differentiable & FloatingPoint
-  )
   func mean(squeezingAxes axes: Int32...) -> Tensor {
     return mean(squeezingAxes: axes)
   }
@@ -1252,10 +1244,6 @@ public extension Tensor where Scalar : Numeric {
   /// - Parameter axes: The dimensions to reduce.
   /// - Precondition: Each value in `axes` must be in the range `-rank...rank`.
   @inlinable @inline(__always)
-  @differentiable(
-    wrt: self, vjp: _vjpSum(squeezingAxes:)
-    where Scalar : Differentiable & FloatingPoint
-  )
   func sum(squeezingAxes axes: [Int32]) -> Tensor {
     return Raw.sum(self, reductionIndices: Tensor<Int32>(axes), keepDims: false)
   }
@@ -1265,10 +1253,6 @@ public extension Tensor where Scalar : Numeric {
   /// - Parameter axes: The dimensions to reduce.
   /// - Precondition: Each value in `axes` must be in the range `-rank...rank`.
   @inlinable @inline(__always)
-  @differentiable(
-    wrt: self, vjp: _vjpSum(squeezingAxes:)
-    where Scalar : Differentiable & FloatingPoint
-  )
   func sum(squeezingAxes axes: Int32...) -> Tensor {
     return sum(squeezingAxes: axes)
   }
@@ -1298,7 +1282,7 @@ public extension Tensor where Scalar : Numeric {
   /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
   @inlinable @inline(__always)
   @differentiable(
-    wrt: self, vjp: _vjpMean(squeezingAxes:)
+    wrt: self, vjp: _vjpMean(alongAxes:)
     where Scalar : Differentiable & FloatingPoint
   )
   func mean(alongAxes axes: [Int32]) -> Tensor {
@@ -1310,10 +1294,7 @@ public extension Tensor where Scalar : Numeric {
   /// - Parameter axes: The dimensions to reduce.
   /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
   @inlinable @inline(__always)
-  @differentiable(
-    wrt: self, vjp: _vjpMean(squeezingAxes:)
-    where Scalar : Differentiable & FloatingPoint
-  )
+  @differentiable(wrt: self where Scalar : Differentiable & FloatingPoint)
   func mean(alongAxes axes: Int32...) -> Tensor {
     return mean(alongAxes: axes)
   }
@@ -1324,7 +1305,7 @@ public extension Tensor where Scalar : Numeric {
   /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
   @inlinable @inline(__always)
   @differentiable(
-    wrt: self, vjp: _vjpSum(squeezingAxes:)
+    wrt: self, vjp: _vjpSum(alongAxes:)
     where Scalar : Differentiable & FloatingPoint
   )
   func sum(alongAxes axes: [Int32]) -> Tensor {
@@ -1336,14 +1317,33 @@ public extension Tensor where Scalar : Numeric {
   /// - Parameter axes: The dimensions to reduce.
   /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
   @inlinable @inline(__always)
-  @differentiable(
-    wrt: self, vjp: _vjpSum(squeezingAxes:)
-    where Scalar : Differentiable & FloatingPoint
-  )
+  @differentiable(wrt: self where Scalar : Differentiable & FloatingPoint)
   func sum(alongAxes axes: Int32...) -> Tensor {
     return sum(alongAxes: axes)
   }
 
+  /// Returns the variance along the specified axes. The reduced dimensions are
+  /// retained with value 1. Does not apply Bessel's correction.
+  /// - Parameter axes: The dimensions to reduce.
+  /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
+  @inlinable @inline(__always)
+  @differentiable(wrt: self where Scalar : Differentiable & FloatingPoint)
+  func variance(alongAxes axes: Int32...) -> Tensor {
+    return variance(alongAxes: axes)
+  }
+
+  /// Returns the variance along the specified axes. The reduced dimensions are
+  /// retained with value 1. Does not apply Bessel's correction.
+  /// - Parameter axes: The dimensions to reduce.
+  /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
+  @inlinable @inline(__always)
+  @differentiable(wrt: self where Scalar : Differentiable & FloatingPoint)
+  func variance(alongAxes axes: [Int32]) -> Tensor {
+    let mean = self.mean(alongAxes: axes)
+    let squaredDiff = (self - mean).squared()
+    return squaredDiff.mean(alongAxes: axes)
+  }
+
   /// Returns the product along the specified axes. The reduced dimensions are
   /// retained with value 1.
   /// - Parameter axes: The dimensions to reduce.

test/AutoDiff/autodiff_function_silgen.swift

@@ -33,10 +33,10 @@ func apply() {
 // CHECK-SILGEN:   [[ORIG:%.*]] = autodiff_function_extract [original] [[DIFFED_COPY]] : $@autodiff @callee_guaranteed (Float) -> Float
 // CHECK-SILGEN:   [[BORROWED_ORIG:%.*]] = begin_borrow [[ORIG]] : $@callee_guaranteed (Float) -> Float
 // CHECK-SILGEN:   apply [[BORROWED_ORIG]]({{%.*}}) : $@callee_guaranteed (Float) -> Float
+// CHECK-SILGEN:   destroy_value [[ORIG]] : $@callee_guaranteed (Float) -> Float
 // CHECK-SILGEN:   [[DIFFED_COPY:%.*]] = copy_value [[DIFFED]] : $@autodiff @callee_guaranteed (Float) -> Float
 // CHECK-SILGEN:   [[ORIG:%.*]] = autodiff_function_extract [original] [[DIFFED_COPY]] : $@autodiff @callee_guaranteed (Float) -> Float
-// CHECK-SILGEN:   [[ORIG_COPY:%.*]] = copy_value [[ORIG]] : $@callee_guaranteed (Float) -> Float
-// CHECK-SILGEN:   return [[ORIG_COPY]] : $@callee_guaranteed (Float) -> Float
+// CHECK-SILGEN:   return [[ORIG]] : $@callee_guaranteed (Float) -> Float
 
 // CHECK-SILGEN-LABEL: @{{.*}}apply{{.*}}
 // CHECK-SILGEN:       [[ORIG:%.*]] = function_ref @{{.*}}thin{{.*}} : $@convention(thin) (Float) -> Float

test/TensorFlowRuntime/tensor_autodiff_runtime.swift

@@ -75,27 +75,27 @@ TensorADTests.testAllBackends("negate") {
 TensorADTests.testAllBackends("sum") {
   let input = Tensor<Float>(randomNormal: [2, 2])
   let sumPullbackScalar = pullback(at: input) { (a: Tensor<Float>) in a.sum() }
-  let sumPullbackSqueezingAxes = pullback(at: input) { (a: Tensor<Float>) in a.sum(squeezingAxes: 0, 1) }
+  // let sumPullbackSqueezingAxes = pullback(at: input) { (a: Tensor<Float>) in a.sum(squeezingAxes: 0, 1) }
   let sumPullbackAlongAxes = pullback(at: input) { (a: Tensor<Float>) in a.sum(alongAxes: 0, 1) }
 
   let expected = Tensor<Float>(ones: [2, 2])
   expectTrue(sumPullbackScalar(Tensor(1)) == expected)
-  expectTrue(sumPullbackSqueezingAxes(Tensor(1)) == expected)
+  // expectTrue(sumPullbackSqueezingAxes(Tensor(1)) == expected)
   expectTrue(sumPullbackAlongAxes(Tensor(1))  == expected)
   expectTrue(sumPullbackScalar(Tensor(3)) == expected * 3)
-  expectTrue(sumPullbackSqueezingAxes(Tensor(3)) == expected * 3)
+  // expectTrue(sumPullbackSqueezingAxes(Tensor(3)) == expected * 3)
   expectTrue(sumPullbackAlongAxes(Tensor(3)) == expected * 3)
 }
 
 TensorADTests.testAllBackends("mean") {
   let meanGradScalar = gradient { (a: Tensor<Float>) in a.mean() }
-  let meanGradSqueezingAxes = gradient { (a: Tensor<Float>) in a.mean(squeezingAxes: 0, 1) }
+  // let meanGradSqueezingAxes = gradient { (a: Tensor<Float>) in a.mean(squeezingAxes: 0, 1) }
   let meanGradAlongAxes = gradient { (a: Tensor<Float>) in a.mean(alongAxes: 0, 1) }
 
   let input = Tensor<Float>(ones: [2, 2])
   let expected = Tensor<Float>(shape: [2, 2], repeating: 0.25)
   expectTrue(meanGradScalar(input) == expected)
-  expectTrue(meanGradSqueezingAxes(input) == expected)
+  // expectTrue(meanGradSqueezingAxes(input) == expected)
   expectTrue(meanGradAlongAxes(input) == expected)
 }
 

utils/update_checkout/update-checkout-config.json

@@ -242,7 +242,7 @@
                 "icu": "release-61-1",
                 "tensorflow": "a6924e6affd935f537cdaf8977094df0e15a7957",
                 "tensorflow-swift-bindings": "10e591340134c37a6c3a1df735a7334a77d5cbc7",
-                "tensorflow-swift-apis": "731ce402ce0bd7459b898b112deb89379bbda893"
+                "tensorflow-swift-apis": "ece4a67ed844919b04d2aec9d131b098aee39413"
             }
         }
     }