@JackCaoG do you know why? cc @miladm

it is not that uncommon that we give slightly different result, is it off by a lot? If is possible that result is different because we run some optimization pass which make code runs faster but not as accurate.

@bdhirsh I hit a problem here. It looks like we do need to redispatch convolution_backward and _convolution to functionalized pass in order to make view ops in their decompositions process be replaced with view_copy ops. However, somehow convolution_backward is not calling into at::native::convolution_backward and thus our own kernel convolution_backward_overrideable is not called.

Here is the dispatcher calls:
convolution:

[call] op=[aten::conv_transpose3d.input], key=[AutogradXLA]
  [call] op=[aten::convolution], key=[AutogradXLA]
   [redispatch] op=[aten::convolution], key=[Functionalize]
    [callBoxed] op=[aten::convolution], key=[XLA]
     [call] op=[aten::_convolution], key=[XLA]
      [redispatchBoxed] op=[aten::_convolution], key=[Meta]
       [call] op=[aten::convolution_overrideable], key=[Functionalize]
        [callBoxed] op=[aten::convolution_overrideable], key=[XLA]
[call] op=[aten::ones_like], key=[Functionalize]

convolution_backward:

[call] op=[aten::convolution_backward], key=[AutogradXLA]
  [redispatch] op=[aten::convolution_backward], key=[Functionalize]
   [callBoxed] op=[aten::convolution_backward], key=[XLA]
    [redispatchBoxed] op=[aten::convolution_backward], key=[Meta]
     [call] op=[aten::new_empty], key=[Functionalize]

Let me know if you need more information.

Is that trace coming from a run on this PR, with these ops removed from XLA? or from before this PR.

In the convolution_backward() trace above, I see aten::convolution_backward called with the Meta key. My guess is that we're calling the meta implementation of convolution_backward for shape inference, either directly in pytorch/xla, or in the functionalization kernel (the fact that the call right before it was with the XLA key makes me think it's coming from the XLA kernel). This will just run shape compute for convolution_backward, so it won't end up dispatching to XLAs implementation - it will run shape compute from core. Is the problem that this is error'ing somehow?

Before this PR, so what the convolution_backward kernal in XLA does is to call

at::functionalization::functionalize_aten_op<ATEN_OP(
      convolution_backward)>::call(grad_output, input, weight, bias_sizes,
                                   stride, padding, dilation, transposed,
                                   output_padding, groups, output_mask);

We did the same thing for _convolution. However, in the Meta kernal of _convolution, it ends up calling convolution_overrideable. But not seeing the same behavior for convolution_backward.

@bdhirsh I applied a hack to make things work.