[functorch] Quick fix for torch.trace backward (pytorch/functorch#527)
* Quick fix for torch.trace backward
Description:
- Quick fix for torch.trace backward
- Updated tests
Note:
We may want to use index_put_ instead of index_put
Fixes pytorch/functorch#501
* Updated the way to decompose trace_backward and added a test
* Update test_vmap.py
