fix stack 2x2 tensor along dimension 1 (#108620)

dotnet · Nov 21, 2024 · d2fdf24 · d2fdf24
1 parent c929990
commit d2fdf24
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/...libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorExtensions.cs b/...libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorExtensions.cs
@@ -3462,7 +3462,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);
         }
@@ -3500,7 +3500,7 @@ 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);
         }

diff --git a/src/libraries/System.Numerics.Tensors/tests/TensorTests.cs b/src/libraries/System.Numerics.Tensors/tests/TensorTests.cs
@@ -1074,6 +1074,27 @@ 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 dimention 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]);
         }
 
         [Fact]