PyTorch Lightning uses deprecated torch_xla function, causing compatibility issues with latest torch_xla versions

[aahila-aws]

Bug description

PyTorch Lightning uses deprecated torch_xla function, causing compatibility issues with latest torch_xla versions.
In this file we see that PyTorch Lightning using an xla function that has been deprecated in the latest versions of torch_xla.

What version are you seeing the problem on?

master

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

AttributeError: module 'torch_xla.runtime' has no attribute 'using_pjrt'

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.4.0):
#- PyTorch Version (e.g., 2.4):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

[byi8220]

Hi, +1 as this issue is affecting me downstream.

Not sure if this is helpful, but it seems the function in question, using_pjrt() was removed in PR pytorch/xla#7787

My understanding (as an outsider to both this codebase and pytorch itself) is that this function was removed in favor of automatic PJRT device detection (see pytorch/xla#7730).

Could one possible quick fix be to follow the same logic the old using_pjrt() function used, and just make a check for device_type() is not None as it does?

I could send out a pretty easy one line PR if that's the case

[lantiga]

Yes I’m planning to fix this for the upcoming 2.5 release, but if you can send a PR that would be super appreciated!

[byi8220]

@lantiga Threw out a proposed fix in #20442

Added 1 new unit test and the relevant tests pass locally (seems CI needs maintainer approval to run), but I am curious, does your CI pipeline cover this case (i.e. does your CI actually install the torch-xla package?)

[lantiga]

Thank you @byi8220! We temporarily disabled TPU checks for PyTorch Lightning but we're going to make it come back soon.