-
Notifications
You must be signed in to change notification settings - Fork 25.7k
Description
As mentioned in #124090, it is possible to end up with a tensor subclass that advertises as contiguous, but secretly has a noncontiguous tensor.
After fixing the FakeTensor problem linked there, we hit another correctness problem: today, AOTAutograd is supposed to be able to always assume that tangents feeding into the backward graph are contiguous. It does so in two steps:
(1) At trace time, when we generate FakeTensor tangents to trace the joint graph, force them to be contiguous
(2) at runtime in the backward, if any of our tangents are not contiguous, force them to be contiguous.
We do this by calling .contiguous(). However, this will incorrectly no-op if our tangent is a contiguous tensor subclass, with a secretly-noncontiguous inner tensor. We should explicitly force the inner tensors to be contiguous as well.
(Technically this will all become unnecessary once we do the work to "trace a pre-dispatch joint graph", although this is still rpetty far away)
cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @anijain2305 @chauhang