ttnn.topk pre-allocated output tensor has incorect output indices

[amalbasaTT]

Describe the bug
ttnn.topk pre-allocated output tensor has incorrect output indices.

To Reproduce
Run unit test provided bellow:

# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.

# SPDX-License-Identifier: Apache-2.0

from loguru import logger
from functools import partial
import pytest
import torch
import ttnn
import traceback

from tests.ttnn.utils_for_testing import comp_equal, check_with_pcc
from tests.tt_eager.python_api_testing.sweep_tests.generation_funcs import gen_func_with_cast_tt
from models.utility_functions import torch_random


def run_topk_output_tests(
    input_shape,
    dim, 
    largest,
    k,
    dtype,
    dlayout,
    in_mem_cfg,
    out_mem_cfg,
    data_seed,
    device,
):
    torch.manual_seed(data_seed)
    # grad tensor
    input_tensor = gen_func_with_cast_tt(
        partial(torch_random, low=-100, high=100, dtype=torch.float32), dtype[0]
    )(input_shape)
    # input tensor 
    
    output_shape = input_shape.copy()
    output_shape[dim] = k
    
    optional_output = [
        gen_func_with_cast_tt(partial(torch_random, low=-100, high=100, dtype=torch.float32), dtype[0])(
            output_shape
        ),
        torch.randint(-1, 1, output_shape, dtype=torch.int16),
    ]

    try:
        # get ref result
        ref_values, ref_indices = torch.topk(
            input_tensor, k, dim=dim, largest=largest, sorted=True
        )

        tt_input_tensor = ttnn.from_torch(
            input_tensor, 
            dtype=dtype[0], 
            layout=dlayout, 
            device=device,
            memory_config=in_mem_cfg)
        tt_output_values = ttnn.from_torch(
            optional_output[0], 
            dtype=dtype[1], 
            layout=dlayout, 
            device=device, 
            memory_config=in_mem_cfg)
        tt_output_indices = ttnn.from_torch(
            optional_output[0], 
            dtype=dtype[1], 
            layout=dlayout, 
            device=device, 
            memory_config=in_mem_cfg)

        ttnn.topk(tt_input_tensor, k=k, dim=dim, largest=largest, sorted=True, out=[tt_output_values, tt_output_indices])
        tt_output_values, tt_output_indices = ttnn.to_torch(tt_output_values), ttnn.to_torch(tt_output_indices).to(torch.int64)
    
    except Exception as e:
        logger.warning(f"Test execution crashed: {e}")
        print(traceback.format_exc())
        raise e

    values_passed, values_comp_output_str = check_with_pcc(ref_values, tt_output_values, 0.999)
    assert values_passed, values_comp_output_str
    indices_paseed, indices_comp_output_str = comp_equal(torch.norm(ref_indices.to(torch.float32)), torch.norm(tt_output_indices.to(torch.float32)))
    assert indices_paseed, f"{indices_comp_output_str}\nGolden Output: {ref_indices}\nActual Output: {tt_output_indices}"
    
    
test_sweep_args = [
    (
        [6, 5, 96, 64],
        -1, 
        True, 
        32,
        [ttnn.bfloat16, ttnn.bfloat16],
        ttnn.TILE_LAYOUT,
        ttnn.DRAM_MEMORY_CONFIG,
        ttnn.DRAM_MEMORY_CONFIG,
        14943539,
    ),
    (
        [2, 24, 222, 256],
        -1, 
        True, 
        32,
        [ttnn.bfloat16, ttnn.bfloat16],
        ttnn.TILE_LAYOUT,
        ttnn.DRAM_MEMORY_CONFIG,
        ttnn.DRAM_MEMORY_CONFIG,
        14943539,
    ),
]


@pytest.mark.parametrize(
    "input_shape, dim, largest, k, dtype, dlayout, in_mem_cfg, out_mem_cfg, data_seed",
    (test_sweep_args),
)
def test_topk_output(input_shape, dim, largest, k, dtype, dlayout, in_mem_cfg, out_mem_cfg, data_seed, device):
    run_topk_output_tests(input_shape, dim, largest, k, dtype, dlayout, in_mem_cfg, out_mem_cfg, data_seed, device)
    

Additional context
For additional info on parameters for which this op is executed, you can check sweep test tests/sweep_framework/sweeps/reduction/topk/topk_output.py.
You can can run it trough elasticsearch by following next steps:

1. Go to tests/sweep_framework/sweeps/reduction/topk/topk_output.py
2. Generate new parameter vectors and run the sweep test

python3 tests/sweep_framework/sweeps_parameter_generator.py --elastic cloud --module-name reduction.topk.topk_output

python3 tests/sweep_framework/sweeps_runner.py --elastic cloud --module-name reduction.topk.topk_output   --suite-name nightly

3. See the error. Results can be found on elastic cloud as explained here: https://github.com/tenstorrent/tt-metal/tree/main/tests/sweep_framework

You can also run it locally by running command:

pytest tests/sweep_framework/sweeps/reduction/topk/topk_output.py -k "test_nightly"