Add context manager to properly convert the precision

[carmocca]

What does this PR do?

The fix in #8208 isn't working as expected. The check for complex64 should be

assert self.complex_buffer.dtype == torch.complex128.

This is because (from https://pytorch.org/docs/stable/generated/torch.set_default_dtype.html):

when the default floating point type is float64 the default complex type is complex128.

However, I didn't find a way to fix it, because moving the model model.double() or model.to(torch.double) does not convert it to complex128 after it has been created.

The only option I see is that the model instantiation is done after torch.set_default_dtype(torch.double) but that's outside of Lightning's scope.
To make this easier, this PR adds the following context manager:

with autodtype(torch.double):
    model = Model()
# or
with trainer.precision_plugin.forward_context():
    model = Model()

This is similar to how #9920 works.

Some other changes included:

• Ensure the default dtype is restored in case of failure.
• Use set_default_dtype instead of set_default_tensor_type.
• Generalize the classes. This is not useful right now because set_default_dtpye only supports floating types currently.

Does your PR introduce any breaking changes? If yes, please list them.

The model is no longer moved.

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

[carmocca]

Not sure whether this is working as expected or a bug. The precision context manager is only active during the forward context, and this hook is not part of it.

Should we instead enter the context manager on setup and exit on teardown?

[tchaton]

I am not sure, but I would say yes. real + img in float32 -> complex64, and real + img in float64 -> complex128. Makes sense to me at least.

[carmocca]

Yes, that's as expected.

The problem here is that we only wrap the precision for the forward hooks.

So, other hooks like setup and on_fit_start are not wrapped and as tested here, they do not use the correct precision.

Maybe we could change this to wrap everything from setup to teardown.

[carmocca]

Note after discussion with Thomas: It's likely we would need to disable it for backward and optimizer.step.

This will also need to be considered for Lite

[tchaton]

Could this be a dataclass ?

[carmocca]

Maybe, but I don't think we want to. It's still a PrecisionPlugin (not a dataclass)

[tchaton]

I am not sure, but I would say yes. real + img in float32 -> complex64, and real + img in float64 -> complex128. Makes sense to me at least.

[carmocca]

This is no longer necessary after the addition of init_module and configure_model