Error when using vmap to compute per-sample-gradient

[YiteWang]

Hi,
When I used vmap to compute per-sample-gradient for the network with following lines which use part of the features as output:
x = x[:,:,::2,::2]
or
y = self.convs(y[:, :, 1:, 1:])

I have encountered the following error:
RuntimeError: vmap: aten::copy_(self, *extra_args) is not possible because there exists a Tensor 'other' in extra_args that has more elements than 'self'. This happened due to 'other' being vmapped over but 'self' not being vmapped over. Please try to use out-of-place operators instead of aten::copy_. If said operator is being called inside the PyTorch framework, please file a bug report instead.

I am wondering if functorch can support such operations.
Thanks!

[Chillee]

@YiteWang Would it be possible to post more of your code? This error (usually) is usually symptomatic of another operator being the problem, which we can solve.

To give more context, as the error message states, vmap is not compatible with certain kinds of in-place operations.

For example

def f(y):
    x = torch.randn(5)
    x.copy_(y)
    return x

vmap(f)(torch.randn(3, 5))

The underlying problem here is that add_ is trying to in-place copy y (a BatchedTensor of shape [5,3]) to a tensor that has fundamentally smaller storage (i.e. 5 elements). If it's out-of-place, then we're allocating a new tensor, and it can be whatever shape we like. But the in-place operation require that the storage not be modified.

However, the somewhat confusing part is that this copy_ may not be occurring in user-code, but instead in PyTorch's library code. So, we need some context to figure out which operations are being called to find the root cause of the issue :)

We also have some work in the pipes to address this issue more broadly. In the above example, it's possible for us to transform the in-place copy into an out-of-place clone. We call this process functionalization, and you can see some work here: #89

TL;DR: The error message you posted is not the underlying issue, could you post more code to help us figure out the real issue :)

[YiteWang]

Thank you for the reply! Sure I can post the code here.
Issue is repeatable as:

import torch
import torch.nn as nn
from functorch import grad, vmap, make_functional_with_buffers, make_functional
batchsize = 2

class testcell(nn.Module):

    def __init__(self):
        super(testcell, self).__init__()
        self.linear = nn.Linear(5, 7)

    def forward(self, inputs):
        feature = self.linear(inputs)
        return feature[1:]

device = torch.cuda.set_device(0)
model = testcell()
model = model.to(device)
input = torch.randn(batchsize, 5).to(device)

func_model, params, buffers = make_functional_with_buffers(model)

def compute_loss(params, buffers, data):
    temp_logits = func_model(params, buffers, data)
    return temp_logits.sum()

print(model(input))

per_sample_grads = vmap(grad(compute_loss), (None, None, 0)) (params, buffers, input)

[Chillee]

Thanks!

The issue is with slice_backward, which causes some problems. There's a couple different ways for us to solve this, so we'll prioritize those.

[Chillee]

@YiteWang This should have been fixed by 68afe85 and pytorch/pytorch#64933.

Can you verify this issue is fixed?

[YiteWang]

@Chillee Yes the issue is fixed.