Use IRootFunctions in Tensor.StdDev (#110641)

* Use IRootFunctions in Tensor.StdDev

* Try fix api compat issue

* Replace Pow+Sum with SumOfSquares

* Drop IPowerFunctions constraint

* Try fix compatability supression

* Fix StdDev stride issue

* Add regression test

* fix test

* Use FlattenedLength in StdDev

* Try add byref to api compat supressions

* Try fix type constraint

---------

Co-authored-by: Tanner Gooding <[email protected]>

Author: lilinus

src/libraries/System.Numerics.Tensors/ref/System.Numerics.Tensors.netcore.cs

@@ -487,7 +487,7 @@ public static void ResizeTo<T>(scoped in System.Numerics.Tensors.Tensor<T> tenso
         public static ref readonly System.Numerics.Tensors.TensorSpan<T> StackAlongDimension<T>(scoped System.ReadOnlySpan<System.Numerics.Tensors.Tensor<T>> tensors, in System.Numerics.Tensors.TensorSpan<T> destination, int dimension) { throw null; }
         public static System.Numerics.Tensors.Tensor<T> Stack<T>(params scoped System.ReadOnlySpan<System.Numerics.Tensors.Tensor<T>> tensors) { throw null; }
         public static ref readonly System.Numerics.Tensors.TensorSpan<T> Stack<T>(scoped in System.ReadOnlySpan<System.Numerics.Tensors.Tensor<T>> tensors, in System.Numerics.Tensors.TensorSpan<T> destination) { throw null; }
-        public static T StdDev<T>(in System.Numerics.Tensors.ReadOnlyTensorSpan<T> x) where T : System.Numerics.IFloatingPoint<T>, System.Numerics.IPowerFunctions<T>, System.Numerics.IAdditionOperators<T, T, T>, System.Numerics.IAdditiveIdentity<T, T> { throw null; }
+        public static T StdDev<T>(in System.Numerics.Tensors.ReadOnlyTensorSpan<T> x) where T : System.Numerics.IFloatingPoint<T>, System.Numerics.IAdditionOperators<T, T, T>, System.Numerics.IAdditiveIdentity<T, T>, System.Numerics.IRootFunctions<T> { throw null; }
         public static System.Numerics.Tensors.Tensor<T> Subtract<T>(in System.Numerics.Tensors.ReadOnlyTensorSpan<T> x, in System.Numerics.Tensors.ReadOnlyTensorSpan<T> y) where T : System.Numerics.ISubtractionOperators<T, T, T> { throw null; }
         public static ref readonly System.Numerics.Tensors.TensorSpan<T> Subtract<T>(scoped in System.Numerics.Tensors.ReadOnlyTensorSpan<T> x, scoped in System.Numerics.Tensors.ReadOnlyTensorSpan<T> y, in System.Numerics.Tensors.TensorSpan<T> destination) where T : System.Numerics.ISubtractionOperators<T, T, T> { throw null; }
         public static System.Numerics.Tensors.Tensor<T> Subtract<T>(in System.Numerics.Tensors.ReadOnlyTensorSpan<T> x, T y) where T : System.Numerics.ISubtractionOperators<T, T, T> { throw null; }

src/libraries/System.Numerics.Tensors/src/CompatibilitySuppressions.xml

@@ -253,4 +253,18 @@
     <Right>lib/net9.0/System.Numerics.Tensors.dll</Right>
     <IsBaselineSuppression>true</IsBaselineSuppression>
   </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0021</DiagnosticId>
+    <Target>M:System.Numerics.Tensors.Tensor.StdDev``1(System.Numerics.Tensors.ReadOnlyTensorSpan{``0}@)``0:T:System.Numerics.IRootFunctions{``0}</Target>
+    <Left>lib/net8.0/System.Numerics.Tensors.dll</Left>
+    <Right>lib/net8.0/System.Numerics.Tensors.dll</Right>
+    <IsBaselineSuppression>true</IsBaselineSuppression>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0021</DiagnosticId>
+    <Target>M:System.Numerics.Tensors.Tensor.StdDev``1(System.Numerics.Tensors.ReadOnlyTensorSpan{``0}@)``0:T:System.Numerics.IRootFunctions{``0}</Target>
+    <Left>lib/net9.0/System.Numerics.Tensors.dll</Left>
+    <Right>lib/net9.0/System.Numerics.Tensors.dll</Right>
+    <IsBaselineSuppression>true</IsBaselineSuppression>
+  </Suppression>
 </Suppressions>

src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorExtensions.cs

@@ -3512,26 +3512,15 @@ public static ref readonly TensorSpan<T> StackAlongDimension<T>(scoped ReadOnlyS
         /// <param name="x">The <see cref="TensorSpan{T}"/> to take the standard deviation of.</param>
         /// <returns><typeparamref name="T"/> representing the standard deviation.</returns>
         public static T StdDev<T>(in ReadOnlyTensorSpan<T> x)
-            where T : IFloatingPoint<T>, IPowerFunctions<T>, IAdditionOperators<T, T, T>, IAdditiveIdentity<T, T>
+            where T : IFloatingPoint<T>, IAdditionOperators<T, T, T>, IAdditiveIdentity<T, T>, IRootFunctions<T>
         {
             T mean = Average(x);
-            Span<T> span = MemoryMarshal.CreateSpan(ref x._reference, (int)x._shape._memoryLength);
-            Span<T> output = new T[x.FlattenedLength];
-            TensorPrimitives.Subtract(span, mean, output);
-            TensorPrimitives.Abs(output, output);
-            TensorPrimitives.Pow((ReadOnlySpan<T>)output, T.CreateChecked(2), output);
-            T sum = TensorPrimitives.Sum((ReadOnlySpan<T>)output);
-            T variance = sum / T.CreateChecked(x._shape._memoryLength);
-
-            if (typeof(T) == typeof(float))
-            {
-                return T.CreateChecked(MathF.Sqrt(float.CreateChecked(variance)));
-            }
-            if (typeof(T) == typeof(double))
-            {
-                return T.CreateChecked(Math.Sqrt(double.CreateChecked(variance)));
-            }
-            return T.Pow(variance, T.CreateChecked(0.5));
+            Tensor<T> temp = CreateUninitialized<T>(x.Lengths);
+            Subtract(x, mean, temp);
+            Abs<T>(temp, temp);
+            T sum = SumOfSquares<T>(temp);
+            T variance = sum / T.CreateChecked(x.FlattenedLength);
+            return T.Sqrt(variance);
         }
         #endregion
 
@@ -6664,6 +6653,19 @@ public static T Sum<T>(scoped in ReadOnlyTensorSpan<T> x)
         }
         #endregion
 
+        #region SumOfSquares
+        /// <summary>
+        /// Sums the squared elements of the specified tensor.
+        /// </summary>
+        /// <param name="x">Tensor to sum squares of</param>
+        /// <returns></returns>
+        internal static T SumOfSquares<T>(scoped in ReadOnlyTensorSpan<T> x)
+            where T : IAdditionOperators<T, T, T>, IAdditiveIdentity<T, T>, IMultiplyOperators<T, T, T>
+        {
+            return TensorPrimitivesHelperSpanInTOut(x, TensorPrimitives.SumOfSquares);
+        }
+        #endregion
+
         #region Tan
         /// <summary>Computes the element-wise tangent of the value in the specified tensor.</summary>
         /// <param name="x">The <see cref="ReadOnlyTensorSpan{T}"/> to take the sin of.</param>

src/libraries/System.Numerics.Tensors/tests/TensorTests.cs

@@ -1124,6 +1124,15 @@ public static void TensorStdDevTests()
             Tensor<float> t0 = Tensor.Create<float>((Enumerable.Range(0, 4).Select(i => (float)i)), [2, 2]);
 
             Assert.Equal(StdDev([0, 1, 2, 3]), Tensor.StdDev<float>(t0), .1);
+
+            // Test that non-contiguous calculations work
+            Tensor<float> fourByFour = Tensor.Create<float>([4, 4]);
+            fourByFour[[0, 0]] = 1f;
+            fourByFour[[0, 1]] = 1f;
+            fourByFour[[1, 0]] = 1f;
+            fourByFour[[1, 1]] = 1f;
+            ReadOnlyTensorSpan<float> upperLeft = fourByFour.AsReadOnlyTensorSpan().Slice([0..2, 0..2]);
+            Assert.Equal(0f, Tensor.StdDev(upperLeft));
         }
 
         public static float StdDev(float[] values)