Lightning-AI/pytorch-lightning · error · ModuleNotFoundError

str(_XLA_AVAILABLE)

Error message

str(_XLA_AVAILABLE)

What it means

Raised by XLACheckpointIO.__init__ when the torch_xla package (XLA support) is not installed. The plugin wrapper for TPU/XLA checkpointing requires torch_xla, and its availability flag (with the import error message) is stringified into this ModuleNotFoundError.

Source

Thrown at src/lightning/fabric/plugins/io/xla.py:40

from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
from lightning.fabric.plugins.io.torch_io import TorchCheckpointIO
from lightning.fabric.utilities.cloud_io import get_filesystem
from lightning.fabric.utilities.types import _PATH

log = logging.getLogger(__name__)


class XLACheckpointIO(TorchCheckpointIO):
    """CheckpointIO that utilizes ``xm.save`` to save checkpoints for TPU training strategies.

    .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.

    """

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        if not _XLA_AVAILABLE:
            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
        super().__init__(*args, **kwargs)

    @override
    def save_checkpoint(self, checkpoint: dict[str, Any], path: _PATH, storage_options: Optional[Any] = None) -> None:
        """Save model/training states as a checkpoint file through state-dump and file-write.

        Args:
            checkpoint: dict containing model and trainer state
            path: write-target path
            storage_options: not used in ``XLACheckpointIO.save_checkpoint``

        Raises:
            TypeError:
                If ``storage_options`` arg is passed in

        """
        if storage_options is not None:
            raise TypeError(

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Install torch_xla appropriate for your platform: `pip install lightning[xla]` or the torch_xla wheel matching your torch version
2. Verify import works: `python -c "import torch_xla"` and fix any underlying ImportError it reports
3. If you don't need TPU/XLA, switch to a different checkpoint IO plugin (TorchCheckpointIO) instead of XLACheckpointIO

Example fix

# before
from lightning.fabric.plugins.io.xla import XLACheckpointIO
io = XLACheckpointIO()  # ModuleNotFoundError

# after
# pip install lightning[xla]
io = XLACheckpointIO()

from lightning.fabric.plugins.io.xla import _XLA_AVAILABLE
if not _XLA_AVAILABLE:
    raise SystemExit("torch_xla not available; install lightning[xla] or use TorchCheckpointIO")

try:
    from lightning.fabric.plugins.io.xla import XLACheckpointIO
except (ModuleNotFoundError, ImportError):
    XLACheckpointIO = None

When it happens

Trigger: Instantiating XLACheckpointIO (or a Fabric/Trainer config that selects the XLA checkpoint IO plugin) in an environment where `import torch_xla` fails or the package is absent.

Common situations: Running on CPU/GPU-only machines, forgetting to install the torch_xla extra (e.g. `pip install lightning[xla]` or torch_xla matching the torch/TPU version), or version mismatch between torch and torch_xla making the import fail.

Understand the failure class

Background: "X is not installed. Please install it with pip install Y": missing optional dependency errors — ImportError/ValueError raised when a library's optional extra was never installed — this error's family across 22 libraries.

Related errors

• str(_XLA_AVAILABLE)
• {str(_XLA_AVAILABLE)}
• {_XLA_AVAILABLE}
• raise ModuleNotFoundError(str(_XLA_AVAILABLE))
• `Trainer.save_checkpoint(..., storage_options=...)` with `st

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/41b13b5e22dce6fc. Report an issue: GitHub.