Type conversions don't seem to preserve pythontensor.

[pytorch]

Type conversions don't seem to preserve pythontensor.

[zou3519]

Fixed by pytorch/pytorch#61458

[zou3519]

The output looks like this:

def forward(self, value_1):
    _tensor_constant0 = self._tensor_constant0
    copy_ = torch.ops.aten.copy_(_tensor_constant0, value_1, False);  _tensor_constant0 = value_1 = None
    _tensor_constant1 = self._tensor_constant1
    return _tensor_constant1

[zou3519]

So the issue is:

• Tensor::to is a composite operation w.r.t. autograd.
• Because it is composite, it gets decomposed before going to the Python key
• The decomposition of Tensor::to creates a new tensor via at::empty: https://github.com/pytorch/pytorch/blob/b45a105acb594cf03e5cc940df7195b0785b8989/aten/src/ATen/native/TensorConversions.cpp#L22-L53
• at::empty doesn't trace back into FX.

Some possible solutions, in no particular order:

1. have at::empty trace back into FX
2. (in pytorch/pytorch) make Tensor::to a primitive w.r.t. autograd. That is, write a backward formula for it. This will make Tensor::to show up in the graph and so whoever is working with the IR doesn't have to deal with an at::empty and then a Tensor::copy_
3. Have functorch have its own set of primitives for autograd. We could declare and define a new operator functorch::to and have DynamicLayerFrontMode always transform aten::to into functorch::to, much like how we're transforming aten::new_zeros into functorch::new_zeros

[zou3519]

There seems to be a general problem that operations that pass through autograd decompose into difficult-to-transform code. #42 is pretty similar.

[Chillee]

Is transforming aten::<x> into functorch::<x> a feasible solution for all these cases?

[zou3519]

If we go down that route we need a way to actually add autograd support to e.g.functorch::to so that it doesn't get decomposed into its constituents. The way that one does autograd support out of tree is by using torch::autograd::Function. There is a TODO item to check that this actually works; if it works then I think it'll be feasible but I'm worried that torch::autograd::Function doesn't compose with the dynamic layer stack.

[zou3519]

pytorch/pytorch#61458 will probably fix this