Xla: `.pow()` produces `aten::`'s in backward.

[taylanbil]

🐛 Bug

.pow(2) produces aten::'s

To Reproduce

Steps to reproduce the behavior:

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


dev = xm.xla_device()
net = nn.Sequential(
    nn.Linear(100,100),
    nn.Linear(100,100),
)
net = net.to(dev)
x = torch.ones(4,100).to(dev)
z = net(x)
z.pow(2).mean().backward()  # produces aten
#z.mean().backward()  # does not produce aten
r = met.metrics_report().split('\n')
for line in r:
    if 'aten::' in line:
        print(line)
xm.mark_step()

OUTPUT:

Counter: aten::conj_out
Counter: aten::view_as_real

Expected behavior

no atens should occur in the metrics report for raising to the power 2.

Environment

• Reproducible on XLA backend [CPU/TPU]: TPU
• torch_xla version: nightly

Additional context

Blocking wav2vec2 and other main models.
Probably the root cause of https://github.com/pytorch/xla/issues/2655

[JackCaoG]

view_as_real is coming from https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/FunctionsManual.cpp#L200 which will indirectly call the view_as_real in https://github.com/pytorch/pytorch/blob/7767dcfc8dd89ed16b97b4915218af4c69985058/aten/src/ATen/native/UnaryOps.cpp#L206 if self is not complex but gradient is complex. Many other backward function share this helper function so it is a common problem.

I think we should implement view_as_real. The functionality is really simple but since it return a writable view, implementation could be tricky.

[taylanbil]

nothing should be complex in the example i gave. everything is a real number already, no?

[JackCaoG]

Another idea is to lower real instead, this way we don't need to deal with writable view generated by view_as_real, @ailzhang What do you think?

[JackCaoG]

@taylanbil out is complex for some reason,

out : [W Copy.cpp:219] Warning: Casting complex values to real discards the imaginary part (function operator())
Columns 1 to 100.001 *
 2.8466 -0.1126  3.1870  0.1087  1.1528 -0.7946 -0.6438  0.2401  0.1829 -0.4525
  2.8466 -0.1126  3.1870  0.1087  1.1528 -0.7946 -0.6438  0.2401  0.1829 -0.4525
  2.8466 -0.1126  3.1870  0.1087  1.1528 -0.7946 -0.6438  0.2401  0.1829 -0.4525
  2.8466 -0.1126  3.1870  0.1087  1.1528 -0.7946 -0.6438  0.2401  0.1829 -0.4525
.
.
.
Columns 91 to 1000.001 *
-0.7061 -0.0488 -0.4167 -1.1050 -0.3238  0.2276  2.5540 -0.0630 -1.2584  1.2156
 -0.7061 -0.0488 -0.4167 -1.1050 -0.3238  0.2276  2.5540 -0.0630 -1.2584  1.2156
 -0.7061 -0.0488 -0.4167 -1.1050 -0.3238  0.2276  2.5540 -0.0630 -1.2584  1.2156
 -0.7061 -0.0488 -0.4167 -1.1050 -0.3238  0.2276  2.5540 -0.0630 -1.2584  1.2156
[ XLAComplexFloatType{4,100} ]

it is calculated by auto out = grad * (exponent * self.pow(exponent - 1)).conj();

[taylanbil]

mathematically, it shouldn't be. perhaps that's where the bug lies?

[JackCaoG]

Let me look into a bit more why out is complex then

[JackCaoG]

Something really weird is going on

 196│ Tensor pow_backward(Tensor grad, const Tensor & self, const Scalar & exponent_) {
 197│   auto exponent = (exponent_.isComplex()) ? exponent_.toComplexDouble() : exponent_.toDouble();
 198├>  if (exponent == 0.0) {
 199│     return at::zeros_like(self, LEGACY_CONTIGUOUS_MEMORY_FORMAT);

[?2004h(gdb) p exponent_.isComplex()
$4 = false
[?2004h(gdb) ptype exponent
type = struct c10::complex<double> [with T = double] {
    T real_;
    T imag_;

even though exponent_ is not complex(and it printed false), code still choose to call toComplexDouble...

[ailzhang]

@anjali411 is working on a fix for this.