[Bug Report] ttnn.mean returns wrong results and shapes

[dmakoviichuk-tt]

Describe the bug Ttnn::mean loses real shape and returns unexpected values Loss before mean: shape: ttnn.Shape([32, 1, 128, 1[32]]) min: 1.5390625 max: 7.59375 mean: 4.6160088 input tensor: shape: ttnn.Shape([32, 1, 128, 1[32]]) min: 1.5390625 max: 7.59375 mean: 4.6160088 mean tensor: shape: ttnn.Shape([1, 1, 32, 32]) min: 0 max: 0.14453125 mean: 0.0001411438 Loss after mean: shape: ttnn.Shape([1, 1, 32, 32]) min: 0 max: 0.14453125 mean: 0.0001411438

To Reproduce

with ttnn.manage_device(device_id=0) as device:
   batch_size = 64
   heads = 4
   num_embeddings = 1
   sentence_length = 32

   torch_tensor = torch.randn(batch_size, heads, sentence_length, num_embeddings, dtype=torch.bfloat16) * 4 + 1
   tt_tensor = ttnn.from_torch(torch_tensor, device=device, layout=ttnn.TILE_LAYOUT)

   ttnn_mean = ttnn.mean(tt_tensor)
   ttnn_mean_torch = ttnn.to_torch(ttnn_mean)
   print(ttnn_mean_torch.shape)
   print(ttnn_mean_torch)
   print(torch.mean(torch_tensor))

Expected behavior ttnn.mean should return mean and right shape.

Additional context Add any other context about the problem here.

[dmakoviichuk-tt]

Hi @bbradelTT it is pretty important for us and we are waiting more then one month for a fix. Could you help us?

[bbradelTT]

I replaced the random values with non random values. E.g. the input is all 5s. What seems to be happening is that the padded shape is used, with padding being 0. Which means that instead of 5/1 it's (5+0+0+..)/32 for shape instead of one value it's one tile

[bbradelTT]

Closing. Note that for larger tensors, the PCC will still be potentially bad until #13621 is fixed.

[rfurko-tt]

@bbradelTT I don't think it's a good reason for ticket to be closed. Please keep it open until it's fixed in main. Thanks in advance!