JIT autograd.grad with gradient_outputs

Hi! I’m trying to use autograd.grad inside a jit decorated function but am having problems with mismatching type hints that I’m unable to resolve. My problem is similar to Error when using TorchScript with autograd · Issue #46483 · pytorch/pytorch · GitHub but I also need to use the grad_outputs argument which is not covered there.

Running the following ( with pytorch 1.9.0 )

def gradient(y, x):
    # grad_outputs = [torch.ones_like(y)]
    grad_outputs = torch.jit.annotate(Optional[Tensor], torch.ones_like(y))
    # grad_outputs = torch.jit.annotate(Optional[List[Optional[Tensor]]], [torch.ones_like(y)]) -> "Expected a List type hint but instead got Optional[List[Optional[Tensor]]]"
    grad = torch.autograd.grad(
        [y], [x], [grad_outputs], create_graph=True, retain_graph=True
    return grad

yields the error


aten::grad(Tensor[] outputs, Tensor[] inputs, Tensor?[]? grad_outputs=None, bool? retain_graph=None, bool create_graph=False, bool allow_unused=False) -> (Tensor?[]):
Expected a value of type 'Optional[List[Optional[Tensor]]]' for argument 'grad_outputs' but instead found type 'List[Tensor]'.
Empty lists default to List[Tensor]. Add a variable annotation to the assignment to create an empty list of another type (torch.jit.annotate(List[T, []]) where T is the type of elements in the list for Python 2)

Any help would be greatly appreciated!


Would you mind creating an issue on GitHub as well, so that we could track and fix it, please?

Thanks for the quick reply. Created the issue here: [JIT] autograd.grad with gradient_outputs argument results in type error · Issue #64728 · pytorch/pytorch · GitHub