vmap support for torch.index_fill

[zou3519]

Tasks

Add vmap support for the following PyTorch operations. That is, each one needs a batching rule.

• torch.index_fill

Expected behavior

Currently, when one does vmap over them, they raise a warning suggesting the batching rule is not implemented:

import torch
from functorch import vmap
x = torch.randn(4, 3, 3)
index = torch.tensor([0, 2])
z = vmap(torch.index_fill, (0, None, None, None))(x, 1, index, -1)
#  UserWarning: There is a performance drop because we have not yet implemented the batching rule

We expect to not see a warning

Read this!

See this note for more context https://github.com/pytorch/pytorch/blob/master/functorch/writing_batching_rules.md

If you're new to developing PyTorch and/or function transforms, we would recommend reading through https://github.com/pytorch/pytorch/wiki/Core-Frontend-Onboarding#unit-8-function-transforms-optional

Should likely be similar to something in https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/functorch/BatchRulesScatterOps.cpp

Testing

pytest test/functorch/test_vmap.py -v -k "op_has_batch_rule and index_fill" should pass (currently, it is an expected failure).

cc @Chillee @samdow @soumith

[qqaatw]

I'm trying to implement this.