Repository: pytorch
Issue: Allow generic python data structure input for torch.autograd.Function
Labels: module: autograd,triaged

### 🚀 The feature, motivation and pitch

I have a custom C++ function that takes in dicts of inputs which can contain different tensors depending on the mode and returns the gradients in the backwards pass. Currently, torch.autograd.Function does not support dict input. It would be nice if it could support dict/list/tuple as input and traverse the input internally and allow the output gradients to be of the same type as the input. There are workarounds for list and tuple, such as exploiting the `*` operator, but not for dict. Here is an example of the desired usage:

```
import torch
from custom_cpp_impl import CustomClass

cls = CustomClass()

class MyCustomFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input_dict):
        ctx.save_for_backward(input_dict)
        output = cls.forward(input_dict)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input_dict = ctx.saved_tensors
        grad_dict = cls.backward(input_dict, grad_output) # Keys would be same as input_dict
        return grad_dict

# Run in "xy" mode
input_dict0 = {'x': torch.tensor(2.0, requires_grad=True),
              'y': torch.tensor(3.0, requires_grad=True)}
output0 = MyCustomFunction.apply(input_dict0)
output0.backward()
print(input_dict0['x'].grad)
print(input_dict0['y'].grad)

# Run in "yz" mode
input_dict1 = {'y': torch.tensor(3.0, requires_grad=True),
              'z': torch.tensor(4.0, requires_grad=True)}
output1 = MyCustomFunction.apply(input_dict1)
output1.backward()
print(input_dict1['y'].grad)
print(input_dict1['z'].grad)
``` 

The returned gradients dict from `backward()` could have the same keys as the input dict, or be prefixed with `'grad_'` for example.

Currently when you try something like this you get `RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn`. I believe there are many users (such as the user in [this](https://discuss.pytorch.org/t/custom-autograd-function-list-dict-input/21862) discussion) who could benefit from this to write more concise custom autograd functions rather than making a different function for each case of the dict input which can be very lengthy and redundant.

### Alternatives

The only alternative I can think of right now is to make a different autograd function for each case of the dict input. For example, if the dict can contain 'x' and 'y' or 'y' and 'z' that get processed in the C++ code differently you would have to do this:
```
import torch
from custom_cpp_impl import CustomClass

cls = CustomClass()

class MyCustomFunctionXY(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, y):
       output = cls.forward({"x": x, "y": y})
       ctx.save_for_backward(x, y)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        x, y = ctx.saved_tensors
        grads = cls.backward({"x": x, "y": y}, grad_output)

        return grads["grad_x"], grads["grad_y"]

class MyCustomFunctionYZ(torch.autograd.Function):
    @staticmethod
    def forward(ctx, y, z):
       output = cls.forward({"y": y, "z": z})
       ctx.save_for_backward(y, z)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        y, z = ctx.saved_tensors
        grads = cls.backward({"y": y, "z": z}, grad_output)

        return grads["grad_y"], grads["grad_z"]

input_dict0 = {'x': torch.tensor(2.0, requires_grad=True),
              'y': torch.tensor(3.0, requires_grad=True)}
input_dict1 = {'y': torch.tensor(3.0, requires_grad=True),
              'z': torch.tensor(4.0, requires_grad=True)}

output0 = MyCustomFunctionXY.apply(input_dict0['x'], input_dict0['y'])
output0.backward()

output1 = MyCustomFunctionYZ.apply(input_dict1['y'], input_dict1['z'])
output1.backward()
```

which is not very concise considering the C++ code is written to take in different combinations of dict input.

### Additional context

_No response_

cc @ezyang @albanD @gqchen @pearu @nikitaved @soulitzer @Varal7 @xmfan