argmax returns incorrect shape

[bbradelTT]

Run pytest

import pytest

import torch

import ttnn
from tests.ttnn.utils_for_testing import assert_with_pcc
from models.utility_functions import torch_random


@pytest.mark.parametrize("batch_size", [1, 32])
@pytest.mark.parametrize("h", [32])
@pytest.mark.parametrize("w", [64000])
@pytest.mark.parametrize("dim", [-1])
def test_argmax(device, batch_size, h, w, dim):

    torch.manual_seed(0)


    torch_input_tensor = torch.randn((1, 1, h, w), dtype=torch.bfloat16)
    torch_output_tensor = torch.argmax(torch_input_tensor, dim=3)
    input_tensor = ttnn.from_torch(torch_input_tensor, device=device)
    output_tensor = ttnn.argmax(input_tensor, dim=3, use_multicore=False if batch_size > 1 else True)
    output_tensor = ttnn.to_torch(output_tensor)
    assert_with_pcc(torch_output_tensor, output_tensor)

Expected:

• Everything passes

Actual:

FAILED test_argmax.py::test_argmax[-1-64000-32-1] - AssertionError: list(expected_pytorch_result.shape)=[1, 1, 32] vs list(actual_pytorch_result.shape)=[1, 1, 1, 32]
FAILED test_argmax.py::test_argmax[-1-64000-32-32] - AssertionError: list(expected_pytorch_result.shape)=[1, 1, 32] vs list(actual_pytorch_result.shape)=[1, 1, 1, 32]

Adding in a reshape fixes the problem, but argmax should set the shapes properly by itself.