Allows functions in `wrap_ad` context to return nested structures

[thomasw21]

torch.autograd.backward only supports list of tensors: https://pytorch.org/docs/stable/generated/torch.autograd.backward.html . This causes the previous assumption to break as self.res_torch could be a nested structure.

We propose to recursively flatten eveything to be passed to the backward step. We also add a test demonstrating that it works.

[Speierers]

Great! Please squash those commits into a single one and then we can merge this PR.

[thomasw21]

I can do this. Though GH usually should have an option to prevent merge commits in settings. I usually have the following deactivated.

[Speierers]

Thanks for your contribution! 🚀