diff --git a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorExtensions.cs b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorExtensions.cs index 6643f7c8a9743..74f570538f34f 100644 --- a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorExtensions.cs +++ b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorExtensions.cs @@ -3456,7 +3456,7 @@ public static Tensor StackAlongDimension(int dimension, params ReadOnlySpa Tensor[] outputs = new Tensor[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(dimension, outputs); } @@ -3494,7 +3494,7 @@ public static ref readonly TensorSpan StackAlongDimension(scoped ReadOnlyS Tensor[] outputs = new Tensor[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(dimension, tensors, destination); } diff --git a/src/libraries/System.Numerics.Tensors/tests/TensorTests.cs b/src/libraries/System.Numerics.Tensors/tests/TensorTests.cs index f1886069e056c..4dde66e651997 100644 --- 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 v1 = Tensor.Create([1, 2, 3, 4], [2, 2]); + Tensor 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]