[fca] support autograd.Function

[xmfan]

Stack from ghstack (oldest at bottom):

• [fca] support accumulate grad hooks #140115
• [fca] support AOT backward #140104
• [aot autograd] functionalize bwd prologue and epilogue #139855
• -> [fca] support autograd.Function #139930
• [fca] support accumulate grad #139121
• [do-not-land] merges #140114
• [WIP] functional autograd + compiled autograd #139098

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @yf225

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/139930

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 60 New Failures

As of commit fe248ef with merge base a951d99 ():

NEW FAILURES - The following jobs have failed:

• Check Labels / Check labels (gh)
  # This PR needs a release notes: label
• Check mergeability of ghstack PR / ghstack-mergeability-check (gh)
  RuntimeError: Command git -C /home/runner/work/pytorch/pytorch cherry-pick -x e2f6fe6 returned non-zero exit code 1
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  Super_SloMo
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_distributed, 1, 1, linux.g5.12xlarge.nvidia.gpu) (gh)
  distributed/_tensor/test_dtensor_compile.py::TestDTensorCompile::test_dtensor_contiguous_dtensor_noncontiguous_local_as_tangent
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  Super_SloMo
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  test_ops_gradients.py::TestBwdGradientsCUDA::test_fn_grad_linalg_vander_cuda_complex128
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  test_modules.py::TestModuleCUDA::test_cpu_gpu_parity_nn_GRU_train_mode_cuda_float32
• inductor / cuda12.1-py3.12-gcc9-sm86 / test (inductor, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  test_ops_gradients.py::TestBwdGradientsCUDA::test_fn_grad_linalg_vander_cuda_complex128
• inductor / cuda12.1-py3.12-gcc9-sm86 / test (inductor, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  test_modules.py::TestModuleCUDA::test_cpu_gpu_parity_nn_GRU_train_mode_cuda_float32
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_avx2, 2, 2, linux.10xlarge.avx2) (gh)
  test_ops_gradients.py::TestBwdGradientsCPU::test_fn_grad_linalg_vander_cpu_complex128
• inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (inductor_avx512, 2, 2, linux.12xlarge) (gh)
  test_ops_gradients.py::TestBwdGradientsCPU::test_fn_grad_linalg_vander_cpu_complex128
• Lint / lintrunner-clang / linux-job (gh)
  >>> Lint for torch/csrc/dynamo/python_compiled_autograd.cpp:
• Lint / lintrunner-noclang / linux-job (gh)
  >>> Lint for torch/_functorch/_aot_autograd/runtime_wrappers.py:
• pull / linux-focal-cuda11.8-py3.10-gcc9 / test (distributed, 1, 3, linux.g4dn.12xlarge.nvidia.gpu) (gh)
  distributed/_tensor/test_dtensor_compile.py::TestDTensorCompile::test_dtensor_contiguous_dtensor_noncontiguous_local_as_tangent
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 1, 5, linux.4xlarge.nvidia.gpu) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_fake_distributed_aot_eager
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh)
  functorch/test_ops.py::TestOperatorsCUDA::test_vjpvjp_matrix_exp_cuda_float32
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 3, 5, linux.4xlarge.nvidia.gpu) (gh)
  inductor/test_compiled_autograd.py::TestCompiledAutograd::test_accumulate_without_zero
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 4, 5, linux.4xlarge.nvidia.gpu) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_custom_tensor_metadata
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 5, 5, linux.4xlarge.nvidia.gpu) (gh)
  functorch/test_ops.py::TestOperatorsCUDA::test_vjpvjp_cumprod_cuda_float32
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 1, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_fake_distributed_aot_eager
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  functorch/test_ops.py::TestOperatorsCUDA::test_vjpvjp_matrix_exp_cuda_float32
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 3, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  functorch/test_ops.py::TestOperatorsCUDA::test_vjpvjp_masked_cumprod_cuda_float32
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 4, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_custom_tensor_metadata
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 5, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  test_nestedtensor.py::TestNestedTensorSubclassCUDA::test_compile_padded_dense_conversion_preserves_metadata_cache_cuda_float32
• pull / linux-focal-py3_9-clang9-xla / test (xla, 1, 1, linux.12xlarge) (gh)
  test_EmbeddingBag_per_sample_weights_and_offsets_xla_int64_int64_float64
• pull / linux-focal-py3.11-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutogradWithCache::test_nested_subclasses_complicated_inps
• pull / linux-focal-py3.11-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_nested_subclasses
• pull / linux-focal-py3.11-clang10 / test (default, 1, 4, linux.4xlarge) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_fake_distributed_aot_eager
• pull / linux-focal-py3.11-clang10 / test (default, 2, 4, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTDispatch::test_aot_dispatch_input_mutation_and_output_alias
• pull / linux-focal-py3.11-clang10 / test (default, 3, 4, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_nested_subclasses
• pull / linux-focal-py3.11-clang10 / test (default, 4, 4, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_custom_tensor_metadata
• pull / linux-focal-py3.11-clang10 / test (dynamo, 1, 3, linux.2xlarge) (gh)
  test_jit.py::TestJit::test_shape_analysis_broadcast
• pull / linux-focal-py3.11-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  nn/test_parametrization.py::TestNNParametrization::test_register_and_remove_parametrization_swap_False
• pull / linux-focal-py3.12-clang10 / test (default, 1, 4, linux.4xlarge) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_fake_distributed_aot_eager
• pull / linux-focal-py3.12-clang10 / test (default, 2, 4, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTDispatch::test_aot_dispatch_input_mutation_and_output_alias
• pull / linux-focal-py3.12-clang10 / test (default, 3, 4, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_nested_subclasses
• pull / linux-focal-py3.12-clang10 / test (default, 4, 4, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_custom_tensor_metadata
• pull / linux-focal-py3.12-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
  test_jit.py::TestJit::test_shape_analysis_broadcast
• pull / linux-focal-py3.12-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  nn/test_parametrization.py::TestNNParametrization::test_register_and_remove_parametrization_swap_False
• pull / linux-focal-py3.9-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutogradWithCache::test_nested_subclasses_complicated_inps
• pull / linux-focal-py3.9-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_nested_subclasses
• pull / linux-focal-py3.9-clang10 / test (default, 1, 4, linux.4xlarge) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_fake_distributed_aot_eager
• pull / linux-focal-py3.9-clang10 / test (default, 2, 4, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTDispatch::test_aot_dispatch_input_mutation_and_output_alias
• pull / linux-focal-py3.9-clang10 / test (default, 3, 4, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_nested_subclasses
• pull / linux-focal-py3.9-clang10 / test (default, 4, 4, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_custom_tensor_metadata
• pull / linux-focal-py3.9-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
  test_jit.py::TestJit::test_shape_analysis_broadcast
• pull / linux-focal-py3.9-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  nn/test_parametrization.py::TestNNParametrization::test_register_and_remove_parametrization_swap_False
• pull / linux-focal-py3.9-clang10-onnx / test (default, 2, 2, linux.2xlarge) (gh)
  onnx/test_pytorch_jit_onnx.py::TestJITIRToONNX_opset14::test_add_sub_with_graph_inputs
• pull / linux-jammy-py3.10-clang15-asan / test (default, 1, 6, linux.4xlarge) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_fake_distributed_aot_eager
• pull / linux-jammy-py3.10-clang15-asan / test (default, 2, 6, linux.4xlarge) (gh)
  inductor/test_compiled_autograd.py::TestCompiledAutograd::test_accumulate_without_zero
• pull / linux-jammy-py3.10-clang15-asan / test (default, 3, 6, linux.4xlarge) (gh)
  inductor/test_compiled_autograd.py::TestCompiledAutograd::test_dynamic_shapes
• pull / linux-jammy-py3.10-clang15-asan / test (default, 4, 6, linux.4xlarge) (gh)
  inductor/test_compiled_autograd.py::TestCompiledAutograd::test_mismatch_fake_tensor_mode
• pull / linux-jammy-py3.10-clang15-asan / test (default, 5, 6, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_nested_subclasses
• pull / linux-jammy-py3.10-clang15-asan / test (default, 6, 6, linux.4xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTModuleSimplified::test_grads_no_force_contiguous_subclass
• pull / linux-jammy-py3.9-gcc11 / test (default, 1, 4, linux.2xlarge) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_fake_distributed_aot_eager
• pull / linux-jammy-py3.9-gcc11 / test (default, 2, 4, linux.2xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTDispatch::test_aot_dispatch_input_mutation_and_output_alias
• pull / linux-jammy-py3.9-gcc11 / test (default, 3, 4, linux.2xlarge) (gh)
  inductor/test_compiled_autograd.py::TestCompiledAutograd::test_accumulate_without_zero
• pull / linux-jammy-py3.9-gcc11 / test (default, 4, 4, linux.2xlarge) (gh)
  functorch/test_aotdispatch.py::TestAOTAutograd::test_custom_tensor_metadata
• pull / linux-jammy-py3.9-gcc11 / test (distributed, 1, 2, linux.2xlarge) (gh)
  distributed/_tensor/test_dtensor_compile.py::TestDTensorCompile::test_dtensor_noncontiguous_output
• pull / linux-jammy-py3.9-gcc11 / test (jit_legacy, 1, 1, linux.2xlarge) (gh)
  test_jit_legacy.py::TestJit::test_shape_analysis_broadcast

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.