StackAlongDimension does not return correct result

[JohnMasen]

Description

StackAlongDimension always use the first parameter to calculate. please see the [Other information] for detailed analyze.

Reproduction Steps

use the code as below

using System.Numerics.Tensors;

#pragma warning disable SYSLIB5001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
int[] f1 = [1, 2, 3, 4];
int[] f2 = [10, 20, 30, 40];

var v1=Tensor.Create(f1, [2,2]);
var v2 = Tensor.Create(f2, [2,2]);
var v3 = StackAlongDimension1(1,[v1,v2]);

Console.WriteLine(v3.ToString(9, 9, 9));

Expected behavior

output should be

[2x2x2], type = System.Int32, isPinned = False
{
{1,2}
{10,20}
{3,4}
{30,40}
}

Actual behavior

output is

[2x2x2], type = System.Int32, isPinned = False
{
{1,2}
{1,2}
{3,4}
{3,4}
}

Regression?

No response

Known Workarounds

No response

Configuration

VS 2022, .NET 8
Windows 11 [Version 10.0.26100.1742]
x64

Other information

by dive into the source code, I found the code is always using the first tensor for stacking.

public static Tensor<T> StackAlongDimension<T>(int dimension, params ReadOnlySpan<Tensor<T>> tensors)
        {
            if (tensors.Length < 2)
                ThrowHelper.ThrowArgument_StackTooFewTensors();

            for (int i = 1; i < tensors.Length; i++)
            {
                if (!TensorHelpers.AreLengthsTheSame<T>(tensors[0], tensors[i]))
                    ThrowHelper.ThrowArgument_StackShapesNotSame();
            }

            if (dimension < 0)
                dimension = tensors[0].Rank - dimension;

            Tensor<T>[] outputs = new Tensor<T>[tensors.Length];
            for (int i = 0; i < tensors.Length; i++)
            {
                outputs[i] = Tensor.Unsqueeze(tensors[0], dimension);    **<--it always use the first tensor**
            }
            return Tensor.ConcatenateOnDimension<T>(dimension, outputs);
        }

I created a PoC version, this version works as exptected.

Tensor<T> StackAlongDimension1<T>(int dimension,Tensor<T>[] tensors)
{

    if (dimension < 0)
    {
        dimension = tensors[0].Rank - dimension;
    }

    Tensor<T>[] array = new Tensor<T>[tensors.Length];
    for (int i = 0; i < tensors.Length; i++)
    {
        array[i] = tensors[i].Unsqueeze(dimension);   **<--use the correct tensor** 
    }
    return Tensor.ConcatenateOnDimension<T>(dimension, array);
}

[dotnet-policy-service]

Tagging subscribers to this area: @dotnet/area-system-numerics-tensors
See info in area-owners.md if you want to be subscribed.

[kasperk81]

@jeffhandley @stephentoub this is a bug in the new API introduced in .NET 9.0. The fix is straightforward #108620. Would it be possible to prioritize this for backporting?

[jeffhandley]

Thanks for the ping to catch my attention, @kasperk81. Yes, once we are able to finish reviewing and merge your fix, we can submit it for backporting. We're likely looking at January's servicing release.