[Feature] Support for mixed TensorDictBase types in torch.autograd.grad

[Xmaster6y]

Description

Support:

inputs = TensorDict(a=torch.randn(2, 3, requires_grad=True))
outputs = inputs + 1
grad_outputs = LazyStackedTensorDict(TensorDict(a=torch.ones(3)), TensorDict(a=torch.ones(3)), stack_dim=0)
grads = torch.autograd.grad(outputs, inputs, grad_outputs)

Motivation and Context

• Mentionned in [Feature] Support for torch.autograd.grad #1417

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[vmoens]

Thanks a mil! Not sure the changes really fix the issue -- see my comment

[Xmaster6y]

What about now? Not sure we can do better. Feel free to close the PR in case of a wontfix.

[Xmaster6y]

@vmoens bump

[vmoens]

Yeah not sure that's exactly what I had in mind.
I think we should avoid things like

tup_outputs = tuple(outputs[k] for k in outputs.keys(True, True))

which are inefficient (we could just use values(...))
But more in general, the idea I had was to check if the various inputs all shared the same type (possibly recursively because you can have a Lazy stack within a regular TD) and if not, specialize the ones that don't have the right type.
It's a bit convoluted though:

• you would need to call td0.apply(check, td1, td2, is_leaf=_is_tensor_collection) where check returns True if all types match
• if check returns False, then you need to cast one of these guys
• then we should implement the logic to do that. I guess something like

  # td: a TD with a LazyStack within
  # plain_td: a TD without lazy stack
   new_td = td.copy().update(plain_td)
  

With that kind of thing we can keep the rest of the function exactly as it it

[Xmaster6y]

Ok I see, thanks for the heads-up.