I want to create a tensor of size m * m * 2, that simulate a square matrix where each cell contain your indexes in a pair. Something like that:

f32[3][3][2]
[
    [
        [0, 0],
        [0, 1],
        [0, 2]
    ],
    [
        [1, 0],
        [1, 1],
        [1, 2]
    ],
    [
        [2, 0],
        [2, 1],
        [2, 2]
    ]
]

I used many Nx functions (iota, broadcast, concatenate, stack, tile, reshape) trying to achieve this result, but the closest tensor I got was the following:

f32[2][3][3]
[
    [
        [0, 0, 0],
        [1, 1, 1],
        [2, 2, 2]
    ],
    [
        [0, 1, 2],
        [0, 1, 2],
        [0, 1, 2]
    ]
]

Any suggestion? Thanks