Unit test passes when ttnn.from_torch is used, but fails when ttnn.Tensor is used

[kmitrovicTT]

Summary

Running tests/ttnn/unit_tests/operations/eltwise/test_binary_composite.py:: test_shape_remainder with shape (1, 1, 3, 3) passes if ttnn.from_torch is used, but fails if ttnn.Tensor is used.

Examples

Passes

@pytest.mark.parametrize(
    "shapes",
    [
        [[1, 1, 3, 3], [1, 1, 3, 3]],
    ],
)
@skip_for_grayskull("Unsupported for Grayskull")
def test_shape_remainder(device, shapes):
    torch.manual_seed(0)
    high = 10
    low = -10

    torch_input_tensor_a = torch.rand(shapes[0], dtype=torch.bfloat16) * (high - low) + low

    high = 9
    low = -9
    torch_input_tensor_b = torch.rand(shapes[1], dtype=torch.bfloat16) * (high - low) + low

    golden_function = ttnn.get_golden_function(ttnn.remainder)
    torch_output_tensor = golden_function(torch_input_tensor_a, torch_input_tensor_b)

    input_tensor_a = ttnn.from_torch(
        torch_input_tensor_a, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG
    )
    input_tensor_b = ttnn.from_torch(
        torch_input_tensor_b, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG
    )

    output_tensor = ttnn.remainder(input_tensor_a, input_tensor_b, memory_config=ttnn.DRAM_MEMORY_CONFIG)
    output_tensor = ttnn.to_torch(output_tensor)

    assert ttnn.pearson_correlation_coefficient(torch_output_tensor, output_tensor) >= 0.999

FAILS with Unsupported shape for tensor conversion from row-major to tile layout...

@pytest.mark.parametrize(
    "shapes",
    [
        [[1, 1, 3, 3], [1, 1, 3, 3]],
    ],
)
@skip_for_grayskull("Unsupported for Grayskull")
def test_shape_remainder(device, shapes):
    torch.manual_seed(0)
    high = 10
    low = -10

    torch_input_tensor_a = torch.rand(shapes[0], dtype=torch.bfloat16) * (high - low) + low

    high = 9
    low = -9
    torch_input_tensor_b = torch.rand(shapes[1], dtype=torch.bfloat16) * (high - low) + low

    golden_function = ttnn.get_golden_function(ttnn.remainder)
    torch_output_tensor = golden_function(torch_input_tensor_a, torch_input_tensor_b)

    input_tensor_a = ttnn.Tensor(torch_input_tensor_a, ttnn.float32).to(ttnn.TILE_LAYOUT).to(device)
    input_tensor_b = ttnn.Tensor(torch_input_tensor_b, ttnn.float32).to(ttnn.TILE_LAYOUT).to(device)

    output_tensor = ttnn.remainder(input_tensor_a, input_tensor_b, memory_config=ttnn.DRAM_MEMORY_CONFIG)
    output_tensor = ttnn.to_torch(output_tensor)

    assert ttnn.pearson_correlation_coefficient(torch_output_tensor, output_tensor) >= 0.999

Repro

pytest -svv tests/ttnn/unit_tests/operations/eltwise/test_binary_composite.py::test_shape_remainder