Lightning-AI/pytorch-lightning · error · ModuleNotFoundError

str(_XLA_AVAILABLE)

Error message

str(_XLA_AVAILABLE)

What it means

XLAPrecision requires torch_xla (PyTorch/XLA for TPU and other XLA devices). The module-level availability check failed and stored the underlying import error in _XLA_AVAILABLE; the constructor converts it into this ModuleNotFoundError when the plugin is created where torch_xla cannot be imported.

Source

Thrown at src/lightning/fabric/plugins/precision/xla.py:41

_PRECISION_INPUT = Literal["32-true", "16-true", "bf16-true"]


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) -> 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

    @override

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Install matching torch_xla: pip install torch_xla with the version corresponding to your torch build (see PyTorch/XLA release compatibility)
2. Confirm 'python -c "import torch_xla"' succeeds in the same environment/interpreter you run Fabric with
3. If you are not targeting TPU/XLA, drop the XLA plugin and use MixedPrecision or the strategy default

Example fix

# before
precision = XLAPrecision("bf16-mixed")  # on a machine without torch_xla
# after (shell)
pip install torch_xla --index-url https://download.pytorch.org/whl/cpu

import importlib.util
if importlib.util.find_spec("torch_xla") is None:
    raise RuntimeError("torch_xla not installed; cannot use XLAPrecision")

import importlib.util

def xla_available() -> bool:
    return importlib.util.find_spec("torch_xla") is not None

try:
    precision = XLAPrecision("bf16-mixed")
except ModuleNotFoundError:
    precision = MixedPrecision("bf16-mixed")

When it happens

Trigger: Instantiating XLAPrecision(...) on a machine where 'import torch_xla' fails — no torch_xla wheel installed, or an xla build that does not match the installed torch version (which typically raises ImportError at import time).

Common situations: Running a TPU-targeted Fabric script locally on CPU/GPU; torch and torch_xia version mismatch after upgrading one but not the other.

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)}
• raise ModuleNotFoundError(str(_XLA_AVAILABLE))
• str(_XLA_AVAILABLE)
• {str(_XLA_AVAILABLE)}
• {_XLA_AVAILABLE}

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