pytorch
ed589dd8 - [functorch] add composition-of-3-transform tests for autograd_function (#90962)

Commit
2 years ago
[functorch] add composition-of-3-transform tests for autograd_function (#90962) This PR adds the following OpInfo tests: - vmap x vjp x vmap - vjp x vmap x vmap - vjp x vjp x vmap These OpInfo tests only run for the autograd_function_db. In general, testing composition of two transforms is sufficient to convince ourselves that functorch works on a given operator. The autograd.Function testing (especially the upcoming generate_vmap_rule) didn't feel rigorous enough to me, so I added these additional tests to convince myself. Test Plan: - new tests Pull Request resolved: https://github.com/pytorch/pytorch/pull/90962 Approved by: https://github.com/samdow, https://github.com/soulitzer
Author
Committer
Parents
Loading