How do I calculate the mean of values with the same label given by a mask for multi-dimensional data?
My input x is a tensor of [B, C, H, W] dim. B is batch size, C number of channels, H height and W width. I have a mask m of [H, W] dim. For each batch size and each channel I want to use the mask m to calculate the mean of all values in [H, W] with the same label.
For example:
B = 2
C = 2
H = 2
W = 3
x = torch.tensor([[[[1,2,3][4,5,6]][[7,8,9][0,1,2]]][[[1,2,3][4,5,6]][[7,8,9][0,1,2]]]])

m = torch.tensor([[0,0,1][0,2,2]])

I expect as output:
mean = torch.tensor([[2.3,3,5.5][5,9,1.5][2.3,3,5.5][5,9,1.5]])

out = torch.tensor([[[[2.3,2.3,3][2.3,5.5,5.5]][[5,5,9][5,1.5,1.5]]][[[2.3,2.3,3][2.3,5.5,5.5]][[5,5,9][5,1.5,1.5]]]])

What I want is the following:
1.) The mean value for each label [mean of label 0, mean of label 1, mean of label 2, mean of label 3, mean of label 4]
2.) I want to replace each value in x with the mean value of its associated label. Output should have the size [B, C, H, W].
I thought to somehow use torch_scatter.scatter_mean.

I figured out how to get the mean values for multi-dimensional data. I flatten H and W and then call scatter_mean from torch_scatter. That gives me the mean values per label for each channel and batch.
x = x.view(B, C, H*W)
m = m.view(B, 1, H*W)
mean = torch_scatter.scatter_mean(x, m, dim=-1)

