Fix DTensor import compatibility for PyTorch < 2.5

[Benoqtr]

What does this PR do?

This PR fixes a compatibility issue related to DTensor import introduced in PyTorch 2.5.

Previously, DTensor was imported only under the condition:

if _torch_distributed_available and is_torch_greater_or_equal("2.5"):
    from torch.distributed.tensor import DTensor

However, this led to a situation where DTensor is not defined at all in PyTorch versions below 2.5. As a result, any later use of:
isinstance(some_tensor, DTensor) would raise a NameError, even if the conditional import was skipped. This PR addresses that issue.

Changes made:

-Added a fallback DTensor = None to ensure the name is always defined.

-Updated downstream code to check if DTensor is not None before using isinstance(..., DTensor).

-Ensures safe and version-compatible handling of DTensor logic across PyTorch versions.

Fixes: N/A (no open issue, but addresses latent compatibility bug)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[vasqu]

Got a few nits / general design choice question

[vasqu]

Is the try catch really necessary?

[vasqu]

Wouldn't it make more sense to create a flag, e.g. _is_dtensor_available = _torch_distributed_available and is_torch_greater_or_equal("2.5") and when an isinstance(..., DTensor) occurs, you simply add the flag check before (_is_dtensor_available and ...).

Not a fan of having an instance check against None.

[Benoqtr]

Thanks @vasqu for the helpful feedback!

I've updated the PR to replace the try/except pattern with a _is_dtensor_available flag as you suggested. Now DTensor is only imported when available, and all isinstance(..., DTensor) checks are guarded with the flag.

Let me know if there's anything else you'd like to adjust.

[vasqu]

Thanks! Could you revert the linter adjustments? Not sure what happened but now there is more than there should be 👀

[Benoqtr]

Thanks! Could you revert the linter adjustments? Not sure what happened but now there is more than there should be 👀

Thanks for the feedback! I'm working on reverting the unintended linter changes now.

I initially applied formatting to try to fix failing code quality checks (which were blocking the CI), but it seems that introduced more changes than necessary. I'll clean it up and make sure only the relevant DTensor logic remains in the PR.

[vasqu]

Thanks for sticking with me :D cc @Cyrilvallez for core maintaner