Lightning-AI/pytorch-lightning · critical · ModuleNotFoundError

{_XLA_AVAILABLE}

Error message

{_XLA_AVAILABLE}

What it means

Raised by XLAPrecision.__init__ when the torch_xla package (and its dependencies) is not importable in the current environment. Lightning's XLA precision plugin delegates all mixed-precision handling to torch_xla, so the plugin cannot be constructed without it. The message string is the import-error text captured by Lightning's module availability check.

Source

Thrown at src/lightning/pytorch/plugins/precision/xla.py:43

from lightning.pytorch.plugins.precision.precision import Precision
from lightning.pytorch.utilities.exceptions import MisconfigurationException


class XLAPrecision(Precision):
    """Plugin for training with XLA.

    Args:
        precision: Full precision (32-true) or half precision (16-true, bf16-true).

    Raises:
        ValueError:
            If unsupported ``precision`` is provided.

    """

    def __init__(self, precision: _PRECISION_INPUT = "32-true") -> None:
        if not _XLA_AVAILABLE:
            raise ModuleNotFoundError(str(_XLA_AVAILABLE))

        supported_precision = get_args(_PRECISION_INPUT)
        if precision not in supported_precision:
            raise ValueError(
                f"`precision={precision!r})` is not supported in XLA."
                f" `precision` must be one of: {supported_precision}."
            )
        self.precision = precision

        if precision == "16-true":
            os.environ["XLA_USE_F16"] = "1"
            self._desired_dtype = torch.float16
        elif precision == "bf16-true":
            os.environ["XLA_USE_BF16"] = "1"
            self._desired_dtype = torch.bfloat16
        else:
            self._desired_dtype = torch.float32

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Install torch_xla matching your PyTorch and Python version (e.g. pip install torch_xla --index-url https://download.pytorch.org/whl/cpu)
  2. If you're not on TPU, switch to a different plugin/strategy (e.g. MixedPrecision for CUDA, no plugin for CPU)
  3. Verify with `python -c "import torch_xla"` to see the underlying import error

Example fix

# before
from lightning.pytorch.plugins import XLAPrecision
plugin = XLAPrecision()  # ModuleNotFoundError on non-TPU machine

# after (CUDA machine)
from lightning.pytorch.plugins import MixedPrecision
plugin = MixedPrecision(precision="16-mixed")
Defensive patterns

Strategy: validation

Validate before calling

from lightning.pytorch.utilities.imports import _XLA_AVAILABLE
if not _XLA_AVAILABLE:
    raise SystemExit("torch_xla not available; use a different precision plugin")

Try / catch

try:
    plugin = XLAPrecision()
except ModuleNotFoundError as e:
    print(f"XLA unavailable ({e}); falling back to MixedPrecision")
    plugin = MixedPrecision(precision="16-mixed")

Prevention

When it happens

Trigger: Constructing XLAPrecision (or passing plugins=XLAPrecision(...) / strategy='xla' with a precision plugin) in an environment where `import torch_xla` fails, e.g. plain CPU/GPU machines or a PyTorch/XLA version mismatch.

Common situations: Running a training script written for TPU/TPOD on a local CUDA or CPU machine; installing pytorch-lightning but forgetting `torch-xla`; upgrading PyTorch to a version with no matching torch_xla wheel.

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


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