please lower std_mean and var_mean

[mobassir94]

torch.std_mean wasn't lowered to xla

here : #2776 @taylanbil shared his minimal test with following code :

import torch.nn as nn
import torch
import torch_xla.core.xla_model as xm
import torch_xla.debug.metrics as met

layer = nn.Linear(100, 100)
d = xm.xla_device()
layer = layer.to(d)
std, mean = torch.std_mean(layer.weight, dim=[1, 0], keepdim=True, unbiased=False)
rep = met.metrics_report()
print(rep)

and observed that the aten counter is there in the metrics report:

...
Counter: XrtSubTuple_Empty
Value: 128
Counter: aten::std_mean
Value: 1
Counter: xla::_copy_from
Value: 2
...

So this op need to be lowered.

var_mean was not lowered either.

but separating works as shown by @taylanbil :

d = xm.xla_device()
layer = layer.to(d)
#std, mean = torch.var_mean(layer.weight, dim=[1, 0], keepdim=True, unbiased=False)
#std, mean = torch.std_mean(layer.weight, dim=[1, 0], keepdim=True, unbiased=False)
std= torch.std(layer.weight, dim=[1, 0], keepdim=True, unbiased=False)
mean = torch.mean(layer.weight.data, dim=[1, 0], keepdim=True)
rep = met.metrics_report()
print(rep)

results in

Counter: xla::mean
Value: 1
Counter: xla::std
Value: 1

[jiasenwu]

Similar issue here in NFNet in timm:
https://github.com/rwightman/pytorch-image-models/blob/779107b693010934ac87c8cecbeb65796e218488/timm/models/layers/std_conv.py#L9

When training on TPU, I see the report of unlowered std_mean. At the moment, we apply a patch to separate the op into std and mean, as suggested above.

[rwightman]

@mobassir94 @jiasenwu to workaround this issue in timm I'll be merging a change that uses F.batch_norm shortly, keeps it all in one kernel (and I can throw in the scale too). Other uses of std_mean for standardization / z-score can also use this approach.

weight = F.batch_norm(
    self.weight.view(1, self.out_channels, -1), None, None,
    weight=(self.gain * self.scale).view(-1),
    training=True, momentum=0., eps=self.eps).reshape_as(self.weight)

[JackCaoG]

We are working on this lowering, should be available on nightly soon.

[miladm]

This feature request is available in the master branch.

[JackCaoG]

Thanks @miladm Close this issue since it is fixed.