TF-404: Add .standardDeviation() to Tensor. (#23763)

This change adds `standardDeviation()` to compute the standard deviation
of a Tensor. It supports optionally specifying the `alongAxes` parameter,
which allows users to reduce only along specified axes.

Author: saeta

stdlib/public/TensorFlow/Ops.swift

@@ -577,6 +577,45 @@ public extension Tensor where Scalar : FloatingPoint & Equatable {
   }
 }
 
+public extension Tensor where Scalar : TensorFlowFloatingPoint {
+  // TODO: standardDeviation() should handle non floating point Tensors.
+
+  /// Returns the standard deviation of the elements 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`.
+  @differentiable(wrt: self)
+  func standardDeviation() -> Tensor {
+    // Reduce along all dimensions.
+    return standardDeviation(alongAxes: Array(0..<shape.rank))
+  }
+
+  /// Returns the standard deviation of the elements 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`.
+  @differentiable(wrt: self)
+  func standardDeviation(alongAxes axes: Int32...) -> Tensor {
+    return standardDeviation(alongAxes: axes)
+  }
+
+  /// Returns the standard deviation of the elements 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)
+  func standardDeviation(alongAxes axes: [Int32]) -> Tensor {
+    return sqrt(variance(alongAxes: axes))
+  }
+}
+
 public extension Tensor where Scalar == Bool {
   /// Computes `!self` element-wise.
   @inlinable @inline(__always)

test/TensorFlowRuntime/tensor.swift

@@ -276,6 +276,26 @@ TensorTests.testAllBackends("SimpleMath") {
                              byError: 0.0001)
 }
 
+TensorTests.testAllBackends("StandardDeviation") {
+  expectEqual(0, Tensor<Float>([1]).standardDeviation().scalarized())
+  expectEqual(
+    0.5,
+    Tensor<Float>([0, 1]).standardDeviation(alongAxes: 0).scalarized())
+  expectEqual(0.5, Tensor<Float>([0, 1]).standardDeviation().scalarized())
+  expectNearlyEqual(
+    2.87228132,
+    Tensor<Float>(rangeFrom: 0, to: 10, stride: 1).standardDeviation().scalarized(),
+    byError: 0.001)
+  let matrix = Tensor<Float>(rangeFrom: 0, to: 10, stride: 1).reshaped(to: [2, 5])
+  expectNearlyEqual(2.87228132,
+                    matrix.standardDeviation().scalarized(),
+                    byError: 0.001)
+  expectPointwiseNearlyEqual(
+    [1.4142, 1.4142],
+    matrix.standardDeviation(alongAxes: 1).array.scalars,
+    byError: 0.001)
+}
+
 TensorTests.testAllBackends("ReductionToScalar") {
   let _: Tensor<Float> = [1, 2, 3, 4, 5]
   // expectEqual(x.mean(), 3)