[Bug Report] Bad shape after transpose in C++

marty1885 commented 4 months ago

Describe the bug

The following C++ code tries to transpose a tensor of shape ttnn.Shape([1, 1, 1[32], 2[32]]). by swapping dimension 0 and 2. Expecting output shape of ttnn.Shape([1, 1, 1[32], 2[32]]). But the actual shape being ttnn.Shape([32, 1, 1, 32]).

#include "tensor/host_buffer/functions.hpp"
#include "tensor/types.hpp"
#include "tt_dnn/op_library/auto_format.hpp"
#include <cstddef>
#include <tt_eager/tensor/tensor.hpp>
#include <ttnn/core.hpp>
#include <ttnn/operations/eltwise/binary/binary.hpp>
#include <ttnn/device.hpp>
#include <tt_dnn/op_library/fully_connected/fully_connected_op.hpp>
#include <tt_dnn/op_library/eltwise_unary/eltwise_unary_op.hpp>
#include <tt_eager/tensor/tensor.hpp>
#include <tt_eager/tt_dnn/op_library/transpose/transpose_op.hpp>

#include "common/bfloat16.hpp"

#include <vector>
#include <iostream>

ttnn::device::Device* device = nullptr;

static tt::tt_metal::Tensor make_random_tensor(tt::tt_metal::Shape s)
{
    static int seed = 42;
     auto b = tt::tt_metal::owned_buffer::create(
        create_random_vector_of_bfloat16_native(
        s[0] * s[1] * s[2] * s[3] * 2
            , 2, seed++, -1));
    tt::tt_metal::Tensor t(OwnedStorage{std::move(b)}, s
        , tt::tt_metal::DataType::BFLOAT16, tt::tt_metal::Layout::ROW_MAJOR);
    return tt::tt_metal::tilize_with_zero_padding(t.to(AutoFormat::GetDefaultDevice()));
}

int main()
{
    device = &ttnn::device::open_device(0);
    AutoFormat::SetDefaultDevice(device);

    auto a = make_random_tensor({1, 1, 1, 2});
    std::cout << "original shape: " << a.shape() << std::endl;
    auto b = tt::tt_metal::transpose(a, 0, 2);
    std::cout << "resulting shaoe: " << b.shape() << std::endl;

    device->close();
}

To Reproduce Steps to reproduce the behavior:

Compile and run the above program
Observe the following output.

                 Device | INFO     | Opening user mode device driver
2024-07-16 10:47:37.745 | INFO     | SiliconDriver   - Detected 1 PCI device : [0]
                  Metal | INFO     | Initializing device 0. Program cache is NOT enabled
                  Metal | INFO     | AI CLK for device 0 is:   1202 MHz
                  Verif | INFO     | Created a random vector of size 2
original shape: ttnn.Shape([1, 1, 1[32], 2[32]])
resulting shaoe: ttnn.Shape([32, 1, 1, 32])
                  Metal | INFO     | Closing device 0
                  Metal | INFO     | Disabling and clearing program cache on device 0
                 Device | INFO     | Closing user mode device drivers

Expected behavior Should result in the correct output shape.

Screenshots If applicable, add screenshots to help explain your problem.

Please complete the following environment information: