[release/9.0-staging] fix stack 2x2 tensor along dimension 1 (#110053)

* fix stack 2x2 tensor along dimension 1

* fixed last stack errors

---------

Co-authored-by: kasperk81 <[email protected]>
Co-authored-by: Michael Sharp <[email protected]>

Author: 3 people

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

@@ -3456,7 +3456,7 @@ public static Tensor<T> StackAlongDimension<T>(int dimension, params ReadOnlySpa
             Tensor<T>[] outputs = new Tensor<T>[tensors.Length];
             for (int i = 0; i < tensors.Length; i++)
             {
-                outputs[i] = Tensor.Unsqueeze(tensors[0], dimension);
+                outputs[i] = Tensor.Unsqueeze(tensors[i], dimension);
             }
             return Tensor.ConcatenateOnDimension<T>(dimension, outputs);
         }
@@ -3494,9 +3494,9 @@ public static ref readonly TensorSpan<T> StackAlongDimension<T>(scoped ReadOnlyS
             Tensor<T>[] outputs = new Tensor<T>[tensors.Length];
             for (int i = 0; i < tensors.Length; i++)
             {
-                outputs[i] = Tensor.Unsqueeze(tensors[0], dimension);
+                outputs[i] = Tensor.Unsqueeze(tensors[i], dimension);
             }
-            return ref Tensor.ConcatenateOnDimension<T>(dimension, tensors, destination);
+            return ref Tensor.ConcatenateOnDimension<T>(dimension, outputs, destination);
         }
         #endregion
 

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

@@ -988,6 +988,7 @@ public static void TensorSetSliceTests()
             Assert.Equal(13, t0[1, 3]);
             Assert.Equal(14, t0[1, 4]);
         }
+
         [Fact]
         public static void TensorStackTests()
         {
@@ -1074,6 +1075,47 @@ public static void TensorStackTests()
             Assert.Equal(8, resultTensor[1, 3, 1]);
             Assert.Equal(9, resultTensor[1, 4, 0]);
             Assert.Equal(9, resultTensor[1, 4, 1]);
+
+            // stacking 2x2 tensors along dimension 1
+            Tensor<int> v1 = Tensor.Create([1, 2, 3, 4], [2, 2]);
+            Tensor<int> v2 = Tensor.Create([10, 20, 30, 40], [2, 2]);
+            
+            resultTensor = Tensor.StackAlongDimension(1, [v1, v2]);
+            
+            Assert.Equal(3, resultTensor.Rank);
+            Assert.Equal(2, resultTensor.Lengths[0]);
+            Assert.Equal(2, resultTensor.Lengths[1]);
+            Assert.Equal(2, resultTensor.Lengths[2]);
+            
+            Assert.Equal(1, resultTensor[0, 0, 0]);
+            Assert.Equal(2, resultTensor[0, 0, 1]);
+            Assert.Equal(10, resultTensor[0, 1, 0]);
+            Assert.Equal(20, resultTensor[0, 1, 1]);
+            
+            Assert.Equal(3, resultTensor[1, 0, 0]);
+            Assert.Equal(4, resultTensor[1, 0, 1]);
+            Assert.Equal(30, resultTensor[1, 1, 0]);
+            Assert.Equal(40, resultTensor[1, 1, 1]);
+
+            resultTensor = Tensor.StackAlongDimension(0, [v1, v2]);
+
+            Tensor<int> resultTensor2 = Tensor.Create<int>([2, 2, 2]);
+            Tensor.StackAlongDimension([v1, v2], resultTensor2, 1);
+
+            Assert.Equal(3, resultTensor2.Rank);
+            Assert.Equal(2, resultTensor2.Lengths[0]);
+            Assert.Equal(2, resultTensor2.Lengths[1]);
+            Assert.Equal(2, resultTensor2.Lengths[2]);
+
+            Assert.Equal(1, resultTensor2[0, 0, 0]);
+            Assert.Equal(2, resultTensor2[0, 0, 1]);
+            Assert.Equal(10, resultTensor2[0, 1, 0]);
+            Assert.Equal(20, resultTensor2[0, 1, 1]);
+
+            Assert.Equal(3, resultTensor2[1, 0, 0]);
+            Assert.Equal(4, resultTensor2[1, 0, 1]);
+            Assert.Equal(30, resultTensor2[1, 1, 0]);
+            Assert.Equal(40, resultTensor2[1, 1, 1]);
         }
 
         [Fact]