Repository: pytorch
Issue: Enable `dict` inputs for `torch.autograd.grad` and `torch.autograd.backward` (usability for `torch.func.functional_call`)
Labels: module: autograd,triaged,module: functorch

### 🚀 The feature, motivation and pitch

`functorch` is deprecated in favor of `torch.func` in PyTorch 2.0. The API [`functorch.make_functional`](https://pytorch.org/functorch/1.13/generated/functorch.make_functional.html#functorch.make_functional) is replaced by [`torch.func.functional_call`](https://pytorch.org/docs/stable/generated/torch.func.functional_call.html#torch.func.functional_call).

- #91811

> #### [`functorch.make_functional`](https://pytorch.org/docs/stable/func.migrating.html#functorch-make-functional)
>
> [torch.func.functional_call()](https://pytorch.org/docs/stable/generated/torch.func.functional_call.html#torch.func.functional_call) is the replacement for [functorch.make_functional](https://pytorch.org/functorch/1.13/generated/functorch.make_functional.html#functorch.make_functional) and [functorch.make_functional_with_buffers](https://pytorch.org/functorch/1.13/generated/functorch.make_functional_with_buffers.html#functorch.make_functional_with_buffers). However, it is not a drop-in replacement.

The `torch.func.functional_call()` API takes `dict[str, Tensor]` inputs for parameters and buffers, while `torch.autograd.{grad,backward}` only supports tensor or tuple of tensors (`Tensor | tuple[Tensor, ...]`) as input. Users need to do manual conversion between `tuple` and `dict`. That is very inconvenient.

This issue requests to support `dict[str, Tensor]` as inputs in `torch.autograd.grad` and `torch.autograd.backward`.

------

Code snippet for example case.

With `fmodel, params = functorch.make_functional(model)`, `params` is `tuple[nn.Parameter, ...]`:

```python
import functorch
import torch
import torch.nn as nn
import torch.nn.functional as F

fmodel: nn.Module
params: tuple[torch.Tensor, ...]
fmodel, params = functorch.make_functional(model)

def calculate_loss(fmodel, params, batch)
    inputs, labels = batch
    outputs = fmodel(params, inputs)
    return F.mse_loss(outputs, labels)

loss: torch.Tensor = calculate_loss(fmodel, params, batch)
grads: tuple[torch.Tensor, ...] = torch.autograd.grad(loss, params)
```

With `torch.func.functional_call(model, params)`, `params` is `dict[str, nn.Parameter]`:

```python
import torch
import torch.nn as nn
import torch.nn.functional as F

params: dict[str, torch.Tensor] = dict(model.named_parameters())

def calculate_loss(model, params, batch)
    inputs, labels = batch
    outputs = torch.func.functional_call(model, params, inputs)
    return F.mse_loss(outputs, labels)

loss: torch.Tensor = calculate_loss(model, params, batch)
grads: tuple[torch.Tensor, ...] = torch.autograd.grad(
    loss,
    tuple(params.values()),  # need to convert `dict` to `tuple` manually
)
grads: dict[str, torch.Tensor] = dict(zip(params.keys(), grads))  # need to convert `tuple` to `dict` manually
```

### Alternatives

_No response_

### Additional context

_No response_

cc @ezyang @albanD @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @zou3519 @Chillee @samdow @kshitij12345 @janeyx99