pytorch

6fea7499 - CompositeImplicitAutograd compliance testing (#65819)

CompositeImplicitAutograd compliance testing (#65819) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/65819 Related to #61669. Functions registered as CompositeImplicitAutograd MUST work for most, if not all, backends. This includes Tensor subclasses. To achieve this, we (PyTorch) impose a set of constraints on how a CompositeImplicitAutograd function can be written. Concretely, this PR adds tests for all OpInfos that checks for compliance. The things that get tested in this PR apply to composite ops and are that: - the op does not change the metadata of a Tensor without performing dispatches - the op does not call set_ or resize_ - the op does not directly access the data ptr The mechanism for the test is to create a new __torch_dispatch__ object, CompositeCompliantTensor. For each operator, we wrap all inputs in CompositeCompliantTensor, turn on python mode for it, and send it through the operator. Non-CompositeImplicitAutograd operators will pass the test because they perform a dispatch to backend code. Here's how CompositeCompliantTensor catches problems: - If it sees set_ or resize_ getting called, it will directly error out - After each operation, CompositeCompliantTensor checks to make sure that its metadata is consistent with that of the thing it is wrapping. If the CompositeImplicitAutograd op modifes the metadata directly (through e.g. the TensorImpl API) then the metadata will go out of sync. - If data_ptr gets called, that returns a nice error (because the storage is meta). CompositeCompliantTensor is written in an interesting way. First off, if a view operation occurs (e.g. `B = A.view_op(...)`), then B.storage() must alias A.storage() where B.storage() is CompositeCompliantTensor's storage, NOT the storage of the tensor it is wrapping. This is an invariant in autograd, see #62182 for details. To handle this we replay the view on A's storage and set it as B's storage. Secondly, there are cases where the metadata is allowed to go out of sync. I believe this is only possible with in-place view functions, like transpose_, t_, squeeze_, unsqueeze_. Those are special cased. Finally, I added a new section to aten/src/ATen/native/README.md about what it means to be CompositeImplicitAutograd Compliant Test Plan: - run tests Reviewed By: ezyang, bdhirsh Differential Revision: D31268369 Pulled By: zou3519 fbshipit-source-id: 31634b1cbe1778ab30196013cfc376ef9bd2e8b1