Enable dict inputs for torch.autograd.grad and torch.autograd.backward (usability for torch.func.functional_call)

[XuehaiPan]

🚀 The feature, motivation and pitch

functorch is deprecated in favor of torch.func in PyTorch 2.0. The API functorch.make_functional is replaced by torch.func.functional_call.

• [torch.func] Add migration guide from functorch #91811

functorch.make_functional

torch.func.functional_call() is the replacement for functorch.make_functional and 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, ...]:

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]:

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

[zou3519]

Leaving this up to @soulitzer (as maintainer of autograd), who will be back next week. It depends on if we (pytorch) are willing to support this (dicts and even pytree support for torch.autograd.grad) going forward.

Let's check in on this in one week.