Add mean and variance ops that take a Tensor<Int32> for axes. (#23864)

Author: bgogul

stdlib/public/TensorFlow/Gradients.swift

@@ -581,11 +581,11 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
   }
 
   @inlinable
-  func _vjpMean(alongAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
+  func _vjpMean(alongAxes axes: Tensor<Int32>) -> (Tensor, (Tensor) -> Tensor) {
     let value = mean(alongAxes: axes)
-    return (value, { [shape = shapeTensor,
-                      count = axes.map { shape[$0] }.reduce(1, *)] in
-      $0.broadcast(toShape: shape) / Tensor(Scalar(count))
+    let count = Raw.gather(params: shapeTensor, indices: axes).product()
+    return (value, { [shape = shapeTensor] in
+      $0.broadcast(toShape: shape) / Tensor(count)
     })
   }
 

stdlib/public/TensorFlow/Ops.swift

@@ -1348,8 +1348,18 @@ public extension Tensor where Scalar : Numeric {
     wrt: self, vjp: _vjpMean(alongAxes:)
     where Scalar : TensorFlowFloatingPoint
   )
+  func mean(alongAxes axes: Tensor<Int32>) -> Tensor {
+    return Raw.mean(self, reductionIndices: axes, keepDims: true)
+  }
+
+  /// Returns the arithmetic mean along the specified axes. The reduced
+  /// dimensions are retained with value 1.
+  /// - 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 : TensorFlowFloatingPoint)
   func mean(alongAxes axes: [Int32]) -> Tensor {
-    return Raw.mean(self, reductionIndices: Tensor<Int32>(axes), keepDims: true)
+    return mean(alongAxes: Tensor<Int32>(axes))
   }
 
   /// Returns the arithmetic mean along the specified axes. The reduced
@@ -1401,12 +1411,22 @@ public extension Tensor where Scalar : Numeric {
   /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
   @inlinable @inline(__always)
   @differentiable(wrt: self where Scalar : TensorFlowFloatingPoint)
-  func variance(alongAxes axes: [Int32]) -> Tensor {
+  func variance(alongAxes axes: Tensor<Int32>) -> Tensor {
     let mean = self.mean(alongAxes: axes)
     let squaredDiff = (self - mean).squared()
     return squaredDiff.mean(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 : TensorFlowFloatingPoint)
+  func variance(alongAxes axes: [Int32]) -> Tensor {
+    return variance(alongAxes: Tensor<Int32>(axes))
+  }
+
   /// Returns the product along the specified axes. The reduced dimensions are
   /// retained with value 1.
   /// - Parameter axes: The dimensions to reduce.