[Bug Report] Operations that reduce broadcasted tensors give incorrect results

[pglusacTT]

Describe the bug
Doing operations that reduce (sum, min, mean...) after having a broadcast gives incorrect results.
If the tensors match in shape the result is correct.

To Reproduce
Run the code below:

import torch
import ttnn

device_id = 0
device = ttnn.open_device(device_id=device_id)

arg1_pt = torch.ones((2, 1))
arg1 = ttnn.from_torch(arg1_pt, dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=device)
arg2_pt = torch.ones((2, 2))
arg2 = ttnn.from_torch(arg2_pt, dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=device)

addition = ttnn.add(arg1, arg2)
    
summ = ttnn.sum(addition)
output_tensor = ttnn.to_torch(summ)


torch_output_tensor = torch.sum(arg1_pt + arg2_pt)

print(output_tensor) # TorchTensor([[68.]])
print(torch_output_tensor) # tensor(8.)

assert torch.allclose(torch_output_tensor, output_tensor, rtol=1e-3), "Error: output_tensor and torch_output_tensor are not close"


ttnn.close_device(device)

Expected behavior
The result of the operations should not be incorrect after broadcasts.

Environment information:

• OS: Ubuntu 22.04
• Version of software: dde5614

[smehtaTT]

@cmaryanTT to prioritize this with @bbradelTT

[bbradelTT]

This may be related to other issues such as #12662
We'll look at this once that issue is unblocked and being worked on.

[vladimirjovanovicTT]

Hi @bbradelTT, could you share a rough ETA for this issue? This is a frequent blocker for Forge training efforts, so an estimate from you will help with our planning.

[bbradelTT]

@vladimirjovanovicTT it turns out that this is related to #12662, which is sufficiently unblocked, and the fix for that issue should fix this issue as well.
Once #16925 is approved and merged, your use case should work.

[vladimirjovanovicTT]

Wow, great news! Thanks.

[bbradelTT]

@vladimirjovanovicTT the change was merged and post commit passes on main. Could you please pull latest main and see if this issue is resolved?

[bbradelTT]

Note that min/max for the first dims would be fixed with a different issue/pr once #16989 is merged.